diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/__init__.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6132a275dc902b75d4387cfa968af1d98498ddd5 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/__init__.py @@ -0,0 +1,51 @@ +# Copyright 2020 The HuggingFace Team. All rights reserved. +# +# 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. +__version__ = "1.12.0" + +from .accelerator import Accelerator +from .big_modeling import ( + cpu_offload, + cpu_offload_with_hook, + disk_offload, + dispatch_model, + init_empty_weights, + init_on_device, + load_checkpoint_and_dispatch, +) +from .data_loader import skip_first_batches +from .inference import prepare_pippy +from .launchers import debug_launcher, notebook_launcher +from .parallelism_config import ParallelismConfig +from .state import PartialState +from .utils import ( + AutocastKwargs, + DataLoaderConfiguration, + DDPCommunicationHookType, + DeepSpeedPlugin, + DistributedDataParallelKwargs, + DistributedType, + FullyShardedDataParallelPlugin, + GradScalerKwargs, + InitProcessGroupKwargs, + ProfileKwargs, + find_executable_batch_size, + infer_auto_device_map, + is_rich_available, + load_checkpoint_in_model, + synchronize_rng_states, +) + + +if is_rich_available(): + from .utils import rich diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/accelerator.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/accelerator.py new file mode 100644 index 0000000000000000000000000000000000000000..5abf98b305afc8a8c51f08ccacf042d173b6a25c --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/accelerator.py @@ -0,0 +1,4350 @@ +# Copyright 2021 The HuggingFace Team. All rights reserved. +# +# 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. + +from __future__ import annotations + +import contextlib +import functools +import json +import math +import os +import re +import shutil +import warnings +from collections import OrderedDict +from contextlib import contextmanager +from functools import partial +from types import MethodType +from typing import Any, Callable, Union + +import torch +import torch.utils.hooks as hooks +from huggingface_hub import split_torch_state_dict_into_shards + +from accelerate.utils.dataclasses import FP8BackendType + +from .big_modeling import _attach_context_parallel_hooks +from .checkpointing import load_accelerator_state, load_custom_state, save_accelerator_state, save_custom_state +from .data_loader import DataLoaderDispatcher, prepare_data_loader, skip_first_batches +from .logging import get_logger +from .optimizer import AcceleratedOptimizer +from .parallelism_config import ParallelismConfig +from .scheduler import AcceleratedScheduler +from .state import AcceleratorState, GradientState, PartialState +from .tracking import LOGGER_TYPE_TO_CLASS, GeneralTracker, filter_trackers +from .utils import ( + MODEL_NAME, + SAFE_WEIGHTS_INDEX_NAME, + SAFE_WEIGHTS_NAME, + SAFE_WEIGHTS_PATTERN_NAME, + WEIGHTS_INDEX_NAME, + WEIGHTS_NAME, + WEIGHTS_PATTERN_NAME, + AORecipeKwargs, + AutocastKwargs, + DataLoaderConfiguration, + DeepSpeedPlugin, + DistributedDataParallelKwargs, + DistributedType, + DynamoBackend, + FP8RecipeKwargs, + FullyShardedDataParallelPlugin, + GradientAccumulationPlugin, + GradScalerKwargs, + InitProcessGroupKwargs, + KwargsHandler, + LoggerType, + MegatronLMPlugin, + MSAMPRecipeKwargs, + PrecisionType, + ProfileKwargs, + ProjectConfiguration, + RNGType, + TERecipeKwargs, + TorchDynamoPlugin, + TorchTensorParallelPlugin, + apply_fp8_autowrap, + check_os_kernel, + clean_state_dict_for_safetensors, + compare_versions, + convert_model, + convert_model_to_fp8_ao, + convert_outputs_to_fp32, + ensure_weights_retied, + extract_model_from_parallel, + fsdp2_apply_ac, + fsdp2_canonicalize_names, + fsdp2_prepare_model, + fsdp2_switch_optimizer_parameters, + gather, + gather_object, + get_fsdp2_grad_scaler, + get_grad_scaler, + get_mixed_precision_context_manager, + get_pretty_name, + has_offloaded_params, + is_bf16_available, + is_bitsandbytes_multi_backend_available, + is_deepspeed_available, + is_ipex_available, + is_lomo_available, + is_megatron_lm_available, + is_mlu_available, + is_msamp_available, + is_musa_available, + is_npu_available, + is_torch_version, + is_torch_xla_available, + is_torchao_available, + is_transformer_engine_available, + is_xpu_available, + load_fsdp_model, + load_fsdp_optimizer, + model_has_dtensor, + pad_across_processes, + parse_choice_from_env, + recursively_apply, + reduce, + release_memory, + save, + save_fsdp_model, + save_fsdp_optimizer, + wait_for_everyone, +) +from .utils.constants import ( + FSDP2_PYTORCH_VERSION, + FSDP_PYTORCH_VERSION, + PROFILE_PATTERN_NAME, + SCALER_NAME, +) +from .utils.modeling import get_state_dict_offloaded_model +from .utils.other import compile_regions, compile_regions_deepspeed, is_compiled_module + + +if is_deepspeed_available(): + from .utils import ( + DeepSpeedEngineWrapper, + DeepSpeedOptimizerWrapper, + DeepSpeedSchedulerWrapper, + DummyOptim, + DummyScheduler, + map_pytorch_optim_to_deepspeed, + ) + +if is_megatron_lm_available(): + from .utils import ( + MegatronEngine, + MegatronLMDummyDataLoader, + MegatronLMDummyScheduler, + MegatronLMOptimizerWrapper, + MegatronLMSchedulerWrapper, + megatron_lm_initialize, + megatron_lm_prepare_data_loader, + megatron_lm_prepare_model_optimizer_scheduler, + ) + +from torch.distributed.algorithms.join import Join + + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + import torch_xla.distributed.xla_multiprocessing as xmp + + +if is_npu_available(check_device=False): + import torch_npu # noqa: F401 + + +try: + from torch.optim.lr_scheduler import LRScheduler +except ImportError: + from torch.optim.lr_scheduler import _LRScheduler as LRScheduler + +logger = get_logger(__name__) + +# Sentinel values for defaults +_split_batches = object() +_dispatch_batches = object() +_even_batches = object() +_use_seedable_sampler = object() + + +class Accelerator: + """ + Creates an instance of an accelerator for distributed training or mixed precision training. + + Args: + device_placement (`bool`, *optional*, defaults to `True`): + Whether or not the accelerator should put objects on device (tensors yielded by the dataloader, model, + etc...). + mixed_precision (`str`, *optional*): + Whether or not to use mixed precision training. Choose from 'no','fp16','bf16' or 'fp8'. Will default to + the value in the environment variable `ACCELERATE_MIXED_PRECISION`, which will use the default value in the + accelerate config of the current system or the flag passed with the `accelerate.launch` command. 'fp8' + requires the installation of transformers-engine. + gradient_accumulation_steps (`int`, *optional*, default to 1): + The number of steps that should pass before gradients are accumulated. A number > 1 should be combined with + `Accelerator.accumulate`. If not passed, will default to the value in the environment variable + `ACCELERATE_GRADIENT_ACCUMULATION_STEPS`. Can also be configured through a `GradientAccumulationPlugin`. + cpu (`bool`, *optional*): + Whether or not to force the script to execute on CPU. Will ignore GPU available if set to `True` and force + the execution on one process only. + dataloader_config (`DataLoaderConfiguration`, *optional*): + A configuration for how the dataloaders should be handled in distributed scenarios. + deepspeed_plugin ([`~utils.DeepSpeedPlugin`] or dict of `str`: [`~utils.DeepSpeedPlugin`], *optional*): + Tweak your DeepSpeed related args using this argument. This argument is optional and can be configured + directly using *accelerate config*. If using multiple plugins, use the configured `key` property of each + plugin to access them from `accelerator.state.get_deepspeed_plugin(key)`. Alias for `deepspeed_plugins`. + fsdp_plugin ([`~utils.FullyShardedDataParallelPlugin`], *optional*): + Tweak your FSDP related args using this argument. This argument is optional and can be configured directly + using *accelerate config* + torch_tp_plugin ([`~utils.TorchTensorParallelPlugin`], *optional*): + Deprecated: use `parallelism_config` with `tp_size` instead. + megatron_lm_plugin ([`~utils.MegatronLMPlugin`], *optional*): + Tweak your MegatronLM related args using this argument. This argument is optional and can be configured + directly using *accelerate config* + rng_types (list of `str` or [`~utils.RNGType`]): + The list of random number generators to synchronize at the beginning of each iteration in your prepared + dataloaders. Should be one or several of: + + - `"torch"`: the base torch random number generator + - `"cuda"`: the CUDA random number generator (GPU only) + - `"xla"`: the XLA random number generator (TPU only) + - `"generator"`: the `torch.Generator` of the sampler (or batch sampler if there is no sampler in your + dataloader) or of the iterable dataset (if it exists) if the underlying dataset is of that type. + + Will default to `["torch"]` for PyTorch versions <=1.5.1 and `["generator"]` for PyTorch versions >= 1.6. + log_with (list of `str`, [`~utils.LoggerType`] or [`~tracking.GeneralTracker`], *optional*): + A list of loggers to be setup for experiment tracking. Should be one or several of: + + - `"all"` + - `"tensorboard"` + - `"wandb"` + - `"trackio"` + - `"aim"` + - `"comet_ml"` + - `"mlflow"` + - `"dvclive"` + - `"swanlab"` + If `"all"` is selected, will pick up all available trackers in the environment and initialize them. Can + also accept implementations of `GeneralTracker` for custom trackers, and can be combined with `"all"`. + project_config ([`~utils.ProjectConfiguration`], *optional*): + A configuration for how saving the state can be handled. + project_dir (`str`, `os.PathLike`, *optional*): + A path to a directory for storing data such as logs of locally-compatible loggers and potentially saved + checkpoints. + step_scheduler_with_optimizer (`bool`, *optional*, defaults to `True`): + Set `True` if the learning rate scheduler is stepped at the same time as the optimizer, `False` if only + done under certain circumstances (at the end of each epoch, for instance). + kwargs_handlers (list of [`~utils.KwargsHandler`], *optional*) + A list of [`~utils.KwargsHandler`] to customize how the objects related to distributed training, profiling + or mixed precision are created. See [kwargs](kwargs) for more information. + dynamo_backend (`str` or [`~utils.DynamoBackend`], *optional*, defaults to `"no"`): + Set to one of the possible dynamo backends to optimize your training with torch dynamo. + dynamo_plugin ([`~utils.TorchDynamoPlugin`], *optional*): + A configuration for how torch dynamo should be handled, if more tweaking than just the `backend` or `mode` + is needed. + gradient_accumulation_plugin ([`~utils.GradientAccumulationPlugin`], *optional*): + A configuration for how gradient accumulation should be handled, if more tweaking than just the + `gradient_accumulation_steps` is needed. + + **Available attributes:** + + - **device** (`torch.device`) -- The device to use. + - **distributed_type** ([`~utils.DistributedType`]) -- The distributed training configuration. + - **local_process_index** (`int`) -- The process index on the current machine. + - **mixed_precision** (`str`) -- The configured mixed precision mode. + - **num_processes** (`int`) -- The total number of processes used for training. + - **optimizer_step_was_skipped** (`bool`) -- Whether or not the optimizer update was skipped (because of + gradient overflow in mixed precision), in which + case the learning rate should not be changed. + - **process_index** (`int`) -- The overall index of the current process among all processes. + - **state** ([`~state.AcceleratorState`]) -- The distributed setup state. + - **sync_gradients** (`bool`) -- Whether the gradients are currently being synced across all processes. + - **use_distributed** (`bool`) -- Whether the current configuration is for distributed training. + """ + + def __init__( + self, + device_placement: bool = True, + split_batches: bool = _split_batches, + mixed_precision: PrecisionType | str | None = None, + gradient_accumulation_steps: int = 1, + cpu: bool = False, + dataloader_config: DataLoaderConfiguration | None = None, + deepspeed_plugin: DeepSpeedPlugin | dict[str, DeepSpeedPlugin] | None = None, + fsdp_plugin: FullyShardedDataParallelPlugin | None = None, + torch_tp_plugin: TorchTensorParallelPlugin | None = None, # Deprecate later, warning in `post_init` + megatron_lm_plugin: MegatronLMPlugin | None = None, + rng_types: list[str | RNGType] | None = None, + log_with: str | LoggerType | GeneralTracker | list[str | LoggerType | GeneralTracker] | None = None, + project_dir: str | os.PathLike | None = None, + project_config: ProjectConfiguration | None = None, + gradient_accumulation_plugin: GradientAccumulationPlugin | None = None, + step_scheduler_with_optimizer: bool = True, + kwargs_handlers: list[KwargsHandler] | None = None, + dynamo_backend: DynamoBackend | str | None = None, + dynamo_plugin: TorchDynamoPlugin | None = None, + deepspeed_plugins: DeepSpeedPlugin | dict[str, DeepSpeedPlugin] | None = None, + parallelism_config: ParallelismConfig | None = None, + ): + self.trackers = [] + if project_config is not None: + self.project_configuration = project_config + else: + self.project_configuration = ProjectConfiguration(project_dir=project_dir) + if project_dir is not None and self.project_dir is None: + self.project_configuration.set_directories(project_dir) + + if mixed_precision is not None: + mixed_precision = str(mixed_precision) + if mixed_precision not in PrecisionType: + raise ValueError( + f"Unknown mixed_precision mode: {mixed_precision}. Choose between {PrecisionType.list()}" + ) + if torch_tp_plugin is not None: + warnings.warn( + "`TorchTensorParallelPlugin` is deprecated and will be removed in a future version of Accelerate. " + "Please use the `ParallelismConfig` with `tp_size` instead.", + FutureWarning, + ) + + if dynamo_plugin is not None and dynamo_backend is not None: + raise ValueError("You cannot pass in both `dynamo_plugin` and `dynamo_backend`, please only pass in one.") + if dynamo_backend is not None: + dynamo_plugin = TorchDynamoPlugin(backend=dynamo_backend) + elif dynamo_plugin is None: + dynamo_plugin = TorchDynamoPlugin() + + if deepspeed_plugins is not None and deepspeed_plugin is not None: + raise ValueError("You cannot pass in both `deepspeed_plugins` and `deepspeed_plugin`.") + elif deepspeed_plugin is not None: + deepspeed_plugins = deepspeed_plugin + + if deepspeed_plugins is None: + # First check if we're creating another `Accelerator` w/o setting `deepspeed_plugin` + if ( + AcceleratorState._shared_state != {} + and AcceleratorState().distributed_type == DistributedType.DEEPSPEED + ): + deepspeed_plugins = AcceleratorState().deepspeed_plugins + else: + # init from env variables + deepspeed_plugins = ( + DeepSpeedPlugin() + if os.environ.get("ACCELERATE_USE_DEEPSPEED", "false").lower() == "true" + else None + ) + else: + # If we're creating a second `Accelerator`, users shouldn't be passing in a `deepspeed_plugin` + if ( + AcceleratorState._shared_state != {} + and AcceleratorState().distributed_type == DistributedType.DEEPSPEED + and AcceleratorState().deepspeed_plugins is not None + ): + raise NotImplementedError( + "You cannot pass in a `deepspeed_plugin` when creating a second `Accelerator`. " + "Please make sure the first `Accelerator` is initialized with all the plugins you want to use." + ) + if isinstance(deepspeed_plugins, dict): + for plugin in deepspeed_plugins.values(): + if not isinstance(plugin, DeepSpeedPlugin): + raise TypeError("`deepspeed_plugin` must be a DeepSpeedPlugin object.") + + if deepspeed_plugins is not None: + os.environ["ACCELERATE_USE_DEEPSPEED"] = "true" # use DeepSpeed if plugin is provided + if not is_deepspeed_available(): + raise ImportError("DeepSpeed is not installed => run `pip install deepspeed` or build it from source.") + if is_mlu_available(): + if compare_versions("deepspeed", "<", "0.15.2"): + raise ImportError("DeepSpeed MLU version must be >= 0.15.2. Please update DeepSpeed.") + elif is_musa_available(): + if compare_versions("deepspeed", "<", "0.14.3"): + raise ImportError("DeepSpeed MUSA version must be >= 0.14.3. Please update DeepSpeed.") + elif compare_versions("deepspeed", "<", "0.9.3"): + raise ImportError("DeepSpeed version must be >= 0.9.3. Please update DeepSpeed.") + + self.deepspeed_engine_wrapped = None + + if os.environ.get("ACCELERATE_USE_FSDP", "false").lower() == "true" or isinstance( + fsdp_plugin, FullyShardedDataParallelPlugin + ): + if not is_torch_version(">=", FSDP_PYTORCH_VERSION): + raise ValueError(f"FSDP requires PyTorch >= {FSDP_PYTORCH_VERSION}") + + if fsdp_plugin is None: # init from env variables + fsdp_plugin = ( + FullyShardedDataParallelPlugin() + if os.environ.get("ACCELERATE_USE_FSDP", "false").lower() == "true" + else None + ) + else: + if not isinstance(fsdp_plugin, FullyShardedDataParallelPlugin): + raise TypeError("`fsdp_plugin` must be a FullyShardedDataParallelPlugin object.") + os.environ["ACCELERATE_USE_FSDP"] = "true" # use FSDP if plugin is provided + + if fsdp_plugin is not None and fsdp_plugin.fsdp_version == 2: + if not is_torch_version(">=", FSDP2_PYTORCH_VERSION): + raise ImportError(f"FSDP2 requires PyTorch >= {FSDP2_PYTORCH_VERSION}") + + if megatron_lm_plugin is None: # init from env variables + megatron_lm_plugin = ( + MegatronLMPlugin() if os.environ.get("ACCELERATE_USE_MEGATRON_LM", "false").lower() == "true" else None + ) + else: + if not isinstance(megatron_lm_plugin, MegatronLMPlugin): + raise TypeError("`megatron_lm_plugin` must be a MegatronLMPlugin object.") + os.environ["ACCELERATE_USE_MEGATRON_LM"] = "true" # use MegatronLM if plugin is provided + + if megatron_lm_plugin: + if not is_megatron_lm_available(): + raise ImportError("Megatron is not installed. please build it from source.") + + # Kwargs handlers + self.ddp_handler = None + self.scaler_handler = None + self.init_handler = None + self.fp8_recipe_handler = None + self.ao_recipe_handler = None + self.te_recipe_handler = None + self.msamp_recipe_handler = None + self.autocast_handler = None + self.profile_handler = None + self.has_lomo_optimizer = False + + found_handlers = set() + handler_class_to_attr = { + DistributedDataParallelKwargs: "ddp_handler", + GradScalerKwargs: "scaler_handler", + InitProcessGroupKwargs: "init_handler", + FP8RecipeKwargs: "fp8_recipe_handler", + AutocastKwargs: "autocast_handler", + ProfileKwargs: "profile_handler", + AORecipeKwargs: "ao_recipe_handler", + TERecipeKwargs: "te_recipe_handler", + MSAMPRecipeKwargs: "msamp_recipe_handler", + } + self.has_fp8_handler = False + if kwargs_handlers is not None: + for handler in kwargs_handlers: + assert isinstance(handler, KwargsHandler), ( + f"Unsupported kwargs handler passed: {handler}, must be one that inherits `accelerate.utils.KwargsHandler`." + ) + # Add the handler class to the set of found handlers + if handler.__class__ in found_handlers: + raise ValueError(f"You can only pass one {handler.__class__} in `kwargs_handlers`.") + found_handlers.add(handler.__class__) + handler_attr = handler_class_to_attr[handler.__class__] + setattr(self, handler_attr, handler) + if "recipe_handler" in handler_attr and not self.has_fp8_handler: + self.has_fp8_handler = True + + if parallelism_config is None: + # TODO: Remove after deprecating tp_plugin + if torch_tp_plugin is not None: + parallelism_config = ParallelismConfig(tp_size=torch_tp_plugin.tp_size) + elif os.environ.get("ACCELERATE_USE_PARALLELISM_CONFIG", "false").lower() == "true": + parallelism_config = ParallelismConfig() + + kwargs = self.init_handler.to_kwargs() if self.init_handler is not None else {} + self.state = AcceleratorState( + mixed_precision=mixed_precision, + cpu=cpu, + dynamo_plugin=dynamo_plugin, + deepspeed_plugin=deepspeed_plugins, + fsdp_plugin=fsdp_plugin, + megatron_lm_plugin=megatron_lm_plugin, + parallelism_config=parallelism_config, + _from_accelerator=True, + **kwargs, + ) + + if self.parallelism_config: + self.state.device_mesh = parallelism_config.get_device_mesh(self.device.type) + self.parallelism_config._validate_accelerator(self) + + self.fp8_enabled = self.state.mixed_precision == "fp8" or mixed_precision == "fp8" + + # Check for automatic FP8 recipe creation + if self.fp8_enabled and not self.has_fp8_handler: + if self.fp8_backend == FP8BackendType.AO: + self.ao_recipe_handler = AORecipeKwargs() + elif self.fp8_backend == FP8BackendType.TE: + self.te_recipe_handler = TERecipeKwargs() + elif self.fp8_backend == FP8BackendType.MSAMP: + self.msamp_recipe_handler = MSAMPRecipeKwargs() + elif self.fp8_backend == FP8BackendType.NO: + # Prioritize AO -> TE -> MSAMP + if is_torchao_available(): + logger.info("Found `torchao` installed, using it for FP8 training.") + self.ao_recipe_handler = AORecipeKwargs() + elif is_transformer_engine_available(): + logger.info("Found `transformer-engine` installed, using it for FP8 training.") + self.te_recipe_handler = TERecipeKwargs() + elif is_msamp_available(): + logger.info("Found `msamp` installed, using it for FP8 training.") + self.msamp_recipe_handler = MSAMPRecipeKwargs() + else: + raise ImportError( + "Tried to train with `fp8` and auto-detect backend, but no FP8-compatible backend was installed. " + "Valid backends are: `torchao`, `transformer-engine`, and `msamp`." + ) + self.has_fp8_handler = True + + self.delayed_fp8_autocast = False + if self.has_fp8_handler: + # We already check if FP8 is available during `self.state` + if not self.fp8_enabled and ( + self.distributed_type not in (DistributedType.FSDP, DistributedType.DEEPSPEED) + ): + raise ValueError("Passing in an FP8 configuration requires setting `mixed_precision='fp8'`.") + self.delayed_fp8_autocast = self.fp8_backend == "TE" and self.distributed_type in ( + DistributedType.MULTI_GPU, + DistributedType.FSDP, + ) + + # TODO: S1ro - this is probably gonna be a problem with other fp8 backends too + if ( + self.fp8_backend == FP8BackendType.AO + and self.state.distributed_type == DistributedType.FSDP + and self.state.fsdp_plugin.cpu_ram_efficient_loading + ): + raise ValueError( + "torchao with FSDP2 and cpu_ram_efficient_loading is not supported, setting `cpu_ram_efficient_loading` to False will fix the issue and work as intended." + ) + + trackers = filter_trackers(log_with, self.logging_dir) + if len(trackers) < 1 and log_with is not None: + warnings.warn(f"`log_with={log_with}` was passed but no supported trackers are currently installed.") + self.log_with = trackers + + if ( + (mixed_precision != "bf16") + and getattr(self.state, "downcast_bfloat", False) + and (self.state.distributedType != DistributedType.XLA) + ): + raise ValueError("Can only use `downcast_bf16` when using `mixed_precision='bf16'` and on a TPU") + + if gradient_accumulation_plugin is not None: + if gradient_accumulation_steps != 1: + raise ValueError( + "You can only pass one of `gradient_accumulation_steps` and `gradient_accumulation_plugin`. Please only pass in the created `GradientAccumulationPlugin` object." + ) + else: + gradient_accumulation_steps = int( + parse_choice_from_env("ACCELERATE_GRADIENT_ACCUMULATION_STEPS", gradient_accumulation_steps) + ) + gradient_accumulation_plugin = GradientAccumulationPlugin(num_steps=gradient_accumulation_steps) + + # If using DeepSpeed, update gradient accumulation steps from the DeepSpeed plugin + self.gradient_state = GradientState( + gradient_accumulation_plugin=gradient_accumulation_plugin, + ) + + self.device_placement = device_placement + if dataloader_config is None: + dataloader_config = DataLoaderConfiguration() + self.dataloader_config = dataloader_config + self.step_scheduler_with_optimizer = step_scheduler_with_optimizer + + # Mixed precision attributes + self.scaler = None + self.native_amp = False + if ( + self.state.mixed_precision == "fp16" + and self.device.type != "cpu" + and self.distributed_type not in (DistributedType.DEEPSPEED, DistributedType.MEGATRON_LM) + ): + self.native_amp = True + supported_device = ("xpu", "cuda", "npu", "xla", "mlu", "musa", "hpu", "sdaa", "mps") + if self.device.type not in supported_device or is_torch_xla_available(check_is_tpu=True): + raise ValueError( + f"fp16 mixed precision requires a device in {supported_device} (not {self.device.type!r})." + ) + if self.device.type == "mps" and not is_torch_version(">=", "2.5.0"): + raise ValueError("fp16 mixed precision with MPS device requires a Pytorch >= 2.5.0") + kwargs = self.scaler_handler.to_kwargs() if self.scaler_handler is not None else {} + + # FSDP2 doesn't use ShardedGradScaler, don't want to modify `get_grad_scaler`, rather create a simple utility + if self.is_fsdp2: + self.scaler = get_fsdp2_grad_scaler(device=self.device.type, **kwargs) + else: + self.scaler = get_grad_scaler(self.distributed_type, **kwargs) + + elif self.state.mixed_precision == "bf16" and self.distributed_type not in ( + DistributedType.DEEPSPEED, + DistributedType.MEGATRON_LM, + ): + if self.device.type in ["cpu", "xpu", "hpu"]: + self.native_amp = True + else: + self.native_amp = is_bf16_available(True) + if not self.native_amp and not is_torch_xla_available(): + raise ValueError("bf16 mixed precision requires PyTorch >= 1.10 and a supported device.") + if self.native_amp and self.device.type == "mps" and not is_torch_version(">=", "2.6.0"): + raise ValueError("bf16 mixed precision with MPS device requires a Pytorch >= 2.6.0") + + # for DeepSpeed, self.state.mixed_precision is always "bf16", + # see https://github.com/huggingface/accelerate/blob/main/src/accelerate/state.py#L968 and + # https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L1263. + elif self.fp8_enabled: + # We always enable `native_amp` for FP8 + self.native_amp = True + if self.fp8_backend == FP8BackendType.MSAMP: + if self.distributed_type == DistributedType.FSDP: + raise NotImplementedError( + "`accelerate` + `MS-AMP` + `FSDP` is not supported at this time. " + "Please consider using deepspeed, which is supported." + ) + elif self.distributed_type != DistributedType.DEEPSPEED: + # MS-AMP requires `GradScaler` even with bf16 autocast w/ single GPU or DDP: + self.scaler = get_grad_scaler(**kwargs) + + # Start of internal step tracking + self.step = 0 + + # Internal references to the training objects + self._optimizers = [] + self._models = [] + self._schedulers = [] + self._dataloaders = [] + self._custom_objects = [] + + # Hooks + self._load_model_state_pre_hook = OrderedDict() + self._save_model_state_pre_hook = OrderedDict() + + # RNG Types + self.rng_types = rng_types + if self.rng_types is None: + self.rng_types = ["generator"] + + # Set a flag tensor for early stopping and other breakpoints + self.flag_tensor = None + + check_os_kernel() + + @property + def deepspeed_plugin(self): + """ + Returns the currently active DeepSpeedPlugin. + + If using multiple plugins, the first one will be the active one by default. Manually call + `accelerator.state.select_deepspeed_plugin(key)` to activate a different plugin. + + If deepspeed is not enabled, this will return `None`. + """ + return self.state.deepspeed_plugin + + @property + def use_distributed(self): + """ + Whether the Accelerator is configured for distributed training + """ + return self.state.use_distributed + + @property + def multi_device(self): + return self.use_distributed and self.distributed_type in ( + DistributedType.MULTI_GPU, + DistributedType.MULTI_MLU, + DistributedType.MULTI_SDAA, + DistributedType.MULTI_MUSA, + DistributedType.MULTI_NPU, + DistributedType.MULTI_XPU, + DistributedType.MULTI_HPU, + ) + + @property + def distributed_type(self): + return self.state.distributed_type + + @property + def num_processes(self): + return self.state.num_processes + + @property + def process_index(self): + return self.state.process_index + + @property + def local_process_index(self): + return self.state.local_process_index + + @property + def device(self): + return self.state.device + + @property + def split_batches(self): + return self.dataloader_config.split_batches + + @property + def dispatch_batches(self): + return self.dataloader_config.dispatch_batches + + @property + def even_batches(self): + return self.dataloader_config.even_batches + + @even_batches.setter + def even_batches(self, value: bool): + self.dataloader_config.even_batches = value + + @property + def use_seedable_sampler(self): + return self.dataloader_config.use_seedable_sampler + + @property + def non_blocking(self): + return self.dataloader_config.non_blocking + + @property + def use_stateful_dataloader(self): + if hasattr(self.dataloader_config, "use_stateful_dataloader"): + return self.dataloader_config.use_stateful_dataloader + return False + + @property + def project_dir(self): + return self.project_configuration.project_dir + + @property + def logging_dir(self): + return self.project_configuration.logging_dir + + @property + def save_iteration(self): + return self.project_configuration.iteration + + @property + def is_main_process(self): + """True for one process only.""" + return self.state.is_main_process + + @property + def is_local_main_process(self): + """True for one process per server.""" + return self.state.is_local_main_process + + @property + def is_last_process(self): + return self.process_index == self.num_processes - 1 + + @property + def mixed_precision(self): + return self.state.mixed_precision + + @property + def is_fsdp2(self): + return self.state.is_fsdp2 + + @property + def is_composable_parallelism_enabled(self): + return self.is_fsdp2 + + @property + def parallelism_config(self) -> Union[ParallelismConfig, None]: + return self.state.parallelism_config + + @property + def torch_device_mesh(self): + return self.state.device_mesh + + @property + def should_save_model(self): + if (pc := self.parallelism_config) is None: + # shouldn't even happen + return self.state.is_local_main_process + _non_model_shard_dims = { + pc.dp_replicate_enabled: "dp_replicate", + pc.cp_enabled: "cp", + } + + # return all( + # self.torch_device_mesh[dim].get_local_rank() == 0 for key, dim in non_model_shard_dims.items() if key + # ) + # TODO: S1ro - this is a temporary solution until we figure out why `save_safe_file` is slow when not all processes + return True + + @property + def tensor_parallel_rank(self) -> int: + """ + Returns the local rank for tensor parallelism. If tensor parallelism is configured but not enabled, returns 0 + since all ranks are assumed to be the same. + """ + if self.parallelism_config: + if self.parallelism_config.tp_enabled: + return self.torch_device_mesh.get_local_rank("tp") + return 0 + raise RuntimeError("Tensor parallelism is not configured. Set `parallelism_config` first.") + + @property + def pipeline_parallel_rank(self) -> int: + """ + Pipeline parallelism is not supported yet. + """ + raise NotImplementedError("Pipeline parallelism is currently not supported in Accelerate.") + + @property + def context_parallel_rank(self) -> int: + """ + Context parallelism is not supported yet. + """ + raise NotImplementedError("Context parallelism is currently not supported in Accelerate.") + + @property + def data_parallel_rank(self) -> int: + """ + Returns the local rank for replicate-based data parallelism. If replicate-based data parallelism is configured + but not enabled, returns 0 since all ranks are assumed to be the same. + """ + if self.parallelism_config: + if self.parallelism_config.dp_replicate_enabled: + return self.torch_device_mesh.get_local_rank("dp_replicate") + return 0 + raise RuntimeError("Data parallelism is not configured. Set `parallelism_config` first.") + + @property + def data_parallel_shard_rank(self) -> int: + """ + Returns the local rank for shard-based data parallelism. If shard-based data parallelism is configured but not + enabled, returns 0 since all ranks are assumed to be the same. + """ + if self.parallelism_config: + if self.parallelism_config.dp_shard_enabled: + return self.torch_device_mesh.get_local_rank("dp_shard") + return 0 + raise RuntimeError("Shard-based data parallelism is not configured. Set `parallelism_config` first.") + + @contextmanager + def split_between_processes(self, inputs: list | tuple | dict | torch.Tensor, apply_padding: bool = False): + """ + Splits `input` between `self.num_processes` quickly and can be then used on that process. Useful when doing + distributed inference, such as with different prompts. + + Note that when using a `dict`, all keys need to have the same number of elements. + + Args: + inputs (`list`, `tuple`, `torch.Tensor`, or `dict` of `list`/`tuple`/`torch.Tensor`): + The input to split between processes. + apply_padding (`bool`, `optional`, defaults to `False`): + Whether to apply padding by repeating the last element of the input so that all processes have the same + number of elements. Useful when trying to perform actions such as `Accelerator.gather()` on the outputs + or passing in less inputs than there are processes. If so, just remember to drop the padded elements + afterwards. + + Example: + + ```python + # Assume there are two processes + from accelerate import Accelerator + + accelerator = Accelerator() + with accelerator.split_between_processes(["A", "B", "C"]) as inputs: + print(inputs) + # Process 0 + ["A", "B"] + # Process 1 + ["C"] + + with accelerator.split_between_processes(["A", "B", "C"], apply_padding=True) as inputs: + print(inputs) + # Process 0 + ["A", "B"] + # Process 1 + ["C", "C"] + ``` + """ + with PartialState().split_between_processes(inputs, apply_padding=apply_padding) as inputs: + yield inputs + + def on_main_process(self, function: Callable[..., Any] | None = None): + """ + A decorator that will run the decorated function on the main process only. Can also be called using the + `PartialState` class. + + Args: + function (`Callable`): The function to decorate. + + Example: + + ```python + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator() + + + >>> @accelerator.on_main_process + ... def print_something(): + ... print("This will be printed by process 0 only.") + + + >>> print_something() + "This will be printed by process 0 only" + ``` + """ + # For times when the `Accelerator` object itself utilizes this decorator. + if function is None: + if "Accelerator." in self.__qualname__: + function = self + else: + raise ValueError( + "The `on_main_process` decorator must be called with a function on an instantiated `Accelerator` object." + ) + + def _inner(*args, **kwargs): + return PartialState().on_main_process(function)(*args, **kwargs) + + return _inner + + def on_local_main_process(self, function: Callable[..., Any] | None = None): + """ + A decorator that will run the decorated function on the local main process only. Can also be called using the + `PartialState` class. + + Args: + function (`Callable`): The function to decorate. + + Example: + ```python + # Assume we have 2 servers with 4 processes each. + from accelerate import Accelerator + + accelerator = Accelerator() + + + @accelerator.on_local_main_process + def print_something(): + print("This will be printed by process 0 only on each server.") + + + print_something() + # On server 1: + "This will be printed by process 0 only" + # On server 2: + "This will be printed by process 0 only" + ``` + """ + # For times when the `Accelerator` object itself utilizes this decorator. + if function is None: + if "Accelerator." in self.__qualname__: + function = self + else: + raise ValueError( + "The `on_local_main_process` decorator must be called with a function on an instantiated `Accelerator` object." + ) + + def _inner(*args, **kwargs): + return PartialState().on_local_main_process(function)(*args, **kwargs) + + return _inner + + def on_last_process(self, function: Callable[..., Any]): + """ + A decorator that will run the decorated function on the last process only. Can also be called using the + `PartialState` class. + + Args: + function (`Callable`): The function to decorate. + + Example: + ```python + # Assume we have 4 processes. + from accelerate import Accelerator + + accelerator = Accelerator() + + + @accelerator.on_last_process + def print_something(): + print(f"Printed on process {accelerator.process_index}") + + + print_something() + "Printed on process 3" + ``` + """ + # For times when the `Accelerator` object itself utilizes this decorator. + if function is None: + if "Accelerator." in self.__qualname__: + function = self + else: + raise ValueError( + "The `on_last_process` decorator must be called with a function on an instantiated `Accelerator` object." + ) + + def _inner(*args, **kwargs): + return PartialState().on_last_process(function)(*args, **kwargs) + + return _inner + + def on_process(self, function: Callable[..., Any] | None = None, process_index: int | None = None): + """ + A decorator that will run the decorated function on a given process index only. Can also be called using the + `PartialState` class. + + Args: + function (`Callable`, `optional`): + The function to decorate. + process_index (`int`, `optional`): + The index of the process on which to run the function. + + Example: + ```python + # Assume we have 4 processes. + from accelerate import Accelerator + + accelerator = Accelerator() + + + @accelerator.on_process(process_index=2) + def print_something(): + print(f"Printed on process {accelerator.process_index}") + + + print_something() + "Printed on process 2" + ``` + """ + # Initial construction of the decorator. + if (self is not None) and (process_index is not None) and (function is None): + return partial(self.on_process, process_index=process_index) + # For times when the `Accelerator` object itself utilizes this decorator. + if function is None: + if "Accelerator." in self.__qualname__: + function = self + else: + raise ValueError( + "The `on_main_process` decorator must be called with a function on an instantiated `Accelerator` object." + ) + + def _inner(*args, **kwargs): + return PartialState().on_process(function, process_index)(*args, **kwargs) + + return _inner + + def on_local_process(self, function: Callable[..., Any] | None = None, local_process_index: int | None = None): + """ + A decorator that will run the decorated function on a given local process index only. Can also be called using + the `PartialState` class. + + Args: + function (`Callable`, *optional*): + The function to decorate. + local_process_index (`int`, *optional*): + The index of the local process on which to run the function. + + Example: + ```python + # Assume we have 2 servers with 4 processes each. + from accelerate import Accelerator + + accelerator = Accelerator() + + + @accelerator.on_local_process(local_process_index=2) + def print_something(): + print(f"Printed on process {accelerator.local_process_index}") + + + print_something() + # On server 1: + "Printed on process 2" + # On server 2: + "Printed on process 2" + ``` + """ + # Initial construction of the decorator. + if (self is not None) and (local_process_index is not None) and (function is None): + return partial(self.on_local_process, local_process_index=local_process_index) + # For times when the `Accelerator` object itself utilizes this decorator. + if function is None: + if "Accelerator." in self.__qualname__: + function = self + else: + raise ValueError( + "The `on_main_process` decorator must be called with a function on an instantiated `Accelerator` object." + ) + + def _inner(*args, **kwargs): + return PartialState().on_local_process(function, local_process_index)(*args, **kwargs) + + return _inner + + @contextmanager + def main_process_first(self): + """ + Lets the main process go first inside a with block. + + The other processes will enter the with block after the main process exits. + + Example: + + ```python + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator() + >>> with accelerator.main_process_first(): + ... # This will be printed first by process 0 then in a seemingly + ... # random order by the other processes. + ... print(f"This will be printed by process {accelerator.process_index}") + ``` + """ + with self.state.main_process_first(): + yield + + @contextmanager + def local_main_process_first(self): + """ + Lets the local main process go inside a with block. + + The other processes will enter the with block after the main process exits. + + Example: + + ```python + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator() + >>> with accelerator.local_main_process_first(): + ... # This will be printed first by local process 0 then in a seemingly + ... # random order by the other processes. + ... print(f"This will be printed by process {accelerator.local_process_index}") + ``` + """ + with self.state.local_main_process_first(): + yield + + @contextmanager + def no_sync(self, model): + """ + A context manager to disable gradient synchronizations across DDP processes by calling + `torch.nn.parallel.DistributedDataParallel.no_sync`. + + If `model` is not in DDP, this context manager does nothing + + Args: + model (`torch.nn.Module`): + PyTorch Module that was prepared with `Accelerator.prepare` + + Example: + + ```python + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator() + >>> dataloader, model, optimizer = accelerator.prepare(dataloader, model, optimizer) + >>> input_a = next(iter(dataloader)) + >>> input_b = next(iter(dataloader)) + + >>> with accelerator.no_sync(): + ... outputs = model(input_a) + ... loss = loss_func(outputs) + ... accelerator.backward(loss) + ... # No synchronization across processes, only accumulate gradients + >>> outputs = model(input_b) + >>> accelerator.backward(loss) + >>> # Synchronization across all processes + >>> optimizer.step() + >>> optimizer.zero_grad() + ``` + """ + if self.is_fsdp2: + model.set_requires_gradient_sync(False) + try: + yield + finally: + model.set_requires_gradient_sync(True) + else: + context = contextlib.nullcontext + if self.use_distributed: + if self.distributed_type != DistributedType.DEEPSPEED or self.state.deepspeed_plugin.zero_stage < 2: + context = getattr(model, "no_sync", context) + + with context(): + yield + + @staticmethod + @contextmanager + def trigger_sync_in_backward(model): + """Trigger the sync of the gradients in the next backward pass of the model after multiple forward passes under + `Accelerator.no_sync` (only applicable in multi-GPU scenarios). + + If the script is not launched in distributed mode, this context manager does nothing. + + Args: + model (`torch.nn.Module`): + The model for which to trigger the gradient synchronization. + + Example: + + ```python + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator() + >>> dataloader, model, optimizer = accelerator.prepare(dataloader, model, optimizer) + + >>> with accelerator.no_sync(): + ... loss_a = loss_func(model(input_a)) # first forward pass + ... loss_b = loss_func(model(input_b)) # second forward pass + >>> accelerator.backward(loss_a) # No synchronization across processes, only accumulate gradients + >>> with accelerator.trigger_sync_in_backward(model): + ... accelerator.backward(loss_b) # Synchronization across all processes + >>> optimizer.step() + >>> optimizer.zero_grad() + ``` + """ + if not isinstance(model, torch.nn.parallel.DistributedDataParallel): + yield + return + + old_require_backward_grad_sync = model.require_backward_grad_sync + old_require_forward_param_sync = model.require_forward_param_sync + + # EXPERIMENTAL: This will force grad sync during `backward()`, but it is unknown if it breaks other DDP features. + # https://github.com/pytorch/pytorch/blob/e1502c0cdbfd17548c612f25d5a65b1e4b86224d/torch/nn/parallel/distributed.py#L1453-L1466 + model.require_backward_grad_sync = True + model.require_forward_param_sync = True + # https://github.com/pytorch/pytorch/blob/e1502c0cdbfd17548c612f25d5a65b1e4b86224d/torch/csrc/distributed/c10d/reducer.cpp#L1371-L1402 + model.reducer.prepare_for_backward([]) + try: + yield + finally: + model.require_backward_grad_sync = old_require_backward_grad_sync + model.require_forward_param_sync = old_require_forward_param_sync + + def _do_sync(self): + "Sets the right `sync_gradients` context and either resets or increases `self.step`" + if self.gradient_state.sync_with_dataloader and self.gradient_state.end_of_dataloader: + self.step = 0 + self.gradient_state._set_sync_gradients(True) + else: + self.step += 1 + self.gradient_state._set_sync_gradients((self.step % self.gradient_state.num_steps) == 0) + + @property + def sync_gradients(self): + return self.gradient_state.sync_gradients + + @sync_gradients.setter + def sync_gradients(self, sync_gradients): + self.gradient_state.sync_gradients = sync_gradients + + @property + def gradient_accumulation_steps(self): + return self.gradient_state.num_steps + + @gradient_accumulation_steps.setter + def gradient_accumulation_steps(self, gradient_accumulation_steps): + self.gradient_state.plugin_kwargs.update({"num_steps": gradient_accumulation_steps}) + + @contextmanager + def accumulate(self, *models): + """ + A context manager that will lightly wrap around and perform gradient accumulation automatically + + Args: + *models (list of `torch.nn.Module`): + PyTorch Modules that were prepared with `Accelerator.prepare`. Models passed to `accumulate()` will + skip gradient syncing during backward pass in distributed training + + Example: + + ```python + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator(gradient_accumulation_steps=1) + >>> dataloader, model, optimizer, scheduler = accelerator.prepare(dataloader, model, optimizer, scheduler) + + >>> for input, output in dataloader: + ... with accelerator.accumulate(model): + ... outputs = model(input) + ... loss = loss_func(outputs) + ... loss.backward() + ... optimizer.step() + ... scheduler.step() + ... optimizer.zero_grad() + ``` + """ + self._do_sync() + + allow_gradient_sync = ( + self.sync_gradients # must sync if sync gradients need to complete an optimizer step + or ( + # the no_sync context stops the gradients from reducing during distributed training + # bringing speedup (potentially at some costs). Here, no_sync can be prevented + # by setting sync_each_batch = True. + self.use_distributed # only relevant in distributed settings + and self.gradient_state.plugin_kwargs.get("sync_each_batch", False) + ) + ) + with contextlib.ExitStack() as cm_stack: + for m in models: + cm_stack.enter_context(contextlib.nullcontext() if allow_gradient_sync else self.no_sync(m)) + yield + + @contextmanager + def join_uneven_inputs(self, joinables, even_batches=None): + """ + A context manager that facilitates distributed training or evaluation on uneven inputs, which acts as a wrapper + around `torch.distributed.algorithms.join`. This is useful when the total batch size does not evenly divide the + length of the dataset. + + Args: + joinables (`list[torch.distributed.algorithms.Joinable]`): + A list of models or optimizers that subclass `torch.distributed.algorithms.Joinable`. Most commonly, a + PyTorch Module that was prepared with `Accelerator.prepare` for DistributedDataParallel training. + even_batches (`bool`, *optional*) + If set, this will override the value of `even_batches` set in the `Accelerator`. If it is not provided, + the default `Accelerator` value wil be used. + + + + `join_uneven_inputs` is only supported for Distributed Data Parallel training on multiple GPUs. For any other + configuration, this method will have no effect. + + + + + + Overriding `even_batches` will not affect iterable-style data loaders. + + + + Example: + + ```python + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator(even_batches=True) + >>> ddp_model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader) + + >>> with accelerator.join_uneven_inputs([ddp_model], even_batches=False): + ... for input, output in dataloader: + ... outputs = model(input) + ... loss = loss_func(outputs) + ... loss.backward() + ... optimizer.step() + ... optimizer.zero_grad() + ``` + """ + if self.multi_device: + dl_even_batches_values = [] + + if even_batches is not None: + iterable_dl_seen = False + # override value in batch sampler for map-style datasets + for dl_idx, dl in enumerate(self._dataloaders): + if isinstance(dl, DataLoaderDispatcher): + iterable_dl_seen = True + continue + dl_even_batches_values.append((dl_idx, dl.batch_sampler.even_batches)) + dl.batch_sampler.even_batches = even_batches + + if iterable_dl_seen: + warnings.warn( + "Overriding even_batches is only supported for map-style datasets, yet some dataloaders given were iterable" + ) + else: + even_batches = self.even_batches + + enable_join = False if even_batches else True + try: + with Join(joinables, enable=enable_join, throw_on_early_termination=False): + yield + finally: + # reset any batch samplers that have been modified + for dl_idx, even_batches_value in dl_even_batches_values: + self._dataloaders[dl_idx].batch_sampler.even_batches = even_batches_value + else: + # Even when disabled, Join expects models to subclass Joinable, so skip entirely for single process runs + if self.distributed_type != DistributedType.NO: + warnings.warn( + "Joining uneven inputs is only supported for multi-GPU training, as a result `join_uneven_inputs` will have no effect." + ) + + with contextlib.nullcontext(joinables): + yield + + def print(self, *args, **kwargs): + """ + Drop in replacement of `print()` to only print once per server. + + Example: + + ```python + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator() + >>> accelerator.print("Hello world!") + ``` + """ + self.state.print(*args, **kwargs) + + def _prepare_one(self, obj, first_pass=False, device_placement=None): + # First pass of preparation: DataLoader, model, optimizer + if first_pass: + if isinstance(obj, torch.utils.data.DataLoader): + return self.prepare_data_loader(obj, device_placement=device_placement) + elif isinstance(obj, torch.nn.Module): + return self.prepare_model(obj, device_placement=device_placement) + elif isinstance(obj, torch.optim.Optimizer): + optimizer = self.prepare_optimizer(obj, device_placement=device_placement) + return optimizer + # Second pass of preparation: LR scheduler (which need the full list of optimizers) + elif isinstance(obj, LRScheduler): + scheduler = self.prepare_scheduler(obj) + return scheduler + # Return the unprocessed object if previous criteria was not met + return obj + + def prepare(self, *args, device_placement=None): + """ + Prepare all objects passed in `args` for distributed training and mixed precision, then return them in the same + order. + + Args: + *args (list of objects): + Any of the following type of objects: + + - `torch.utils.data.DataLoader`: PyTorch Dataloader + - `torch.nn.Module`: PyTorch Module + - `torch.optim.Optimizer`: PyTorch Optimizer + - `torch.optim.lr_scheduler.LRScheduler`: PyTorch LR Scheduler + + device_placement (`list[bool]`, *optional*): + Used to customize whether automatic device placement should be performed for each object passed. Needs + to be a list of the same length as `args`. Not compatible with DeepSpeed or FSDP. + + + + You don't need to prepare a model if you only use it for inference without any kind of mixed precision + + + + Examples: + + ```python + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator() + >>> # Assume a model, optimizer, data_loader and scheduler are defined + >>> model, optimizer, data_loader, scheduler = accelerator.prepare(model, optimizer, data_loader, scheduler) + ``` + + ```python + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator() + >>> # Assume a model, optimizer, data_loader and scheduler are defined + >>> device_placement = [True, True, False, False] + >>> # Will place the first two items passed in automatically to the right device but not the last two. + >>> model, optimizer, data_loader, scheduler = accelerator.prepare( + ... model, optimizer, data_loader, scheduler, device_placement=device_placement + ... ) + ``` + """ + if device_placement is None: + device_placement = [None for _ in args] + elif self.distributed_type in (DistributedType.DEEPSPEED, DistributedType.MEGATRON_LM): + raise ValueError("You can't customize device placements with DeepSpeed or Megatron-LM.") + elif len(device_placement) != len(args): + raise ValueError( + f"`device_placement` should be a list with {len(args)} elements (the number of objects passed)." + ) + + for obj in args: + # TODO: Look at enabling native TP training directly with a proper config + if ( + isinstance(obj, torch.nn.Module) + and self.verify_device_map(obj) + and self.distributed_type != DistributedType.NO + and os.environ.get("ACCELERATE_BYPASS_DEVICE_MAP", "false") != "true" + ): + raise ValueError( + "You can't train a model that has been loaded with `device_map='auto'` in any distributed mode." + " Please rerun your script specifying `--num_processes=1` or by launching with `python {{myscript.py}}`." + ) + + if self.distributed_type == DistributedType.DEEPSPEED: + model_count = 0 + for obj in args: + if isinstance(obj, torch.nn.Module): + model_count += 1 + if model_count > 1: + raise AssertionError( + "You can't use same `Accelerator()` instance with multiple models when using DeepSpeed" + ) + + # On TPUs, putting the model on the XLA device will create new parameters, so the corresponding optimizer will + # have parameters disconnected from the model (so no training :-( ). + # If the model and optimizer have parameters on different devices we raise an error. + if self.distributed_type == DistributedType.XLA: + model_device, optimizer_device = self._get_devices() + if model_device is not None and optimizer_device is not None and model_device != optimizer_device: + raise ValueError( + "The model and the optimizer parameters are not on the same device, which probably means you " + "created an optimizer around your model **before** putting on the device. Make sure the line " + "model.to(device) is before the optimizer creation in your script or remove it entirely and use " + "the flag default value for `device_placement` in your `Accelerator` to let it handle that " + "part for you." + ) + + if self.is_fsdp2: + model_count = 0 + optimizer_count = 0 + for i, obj in enumerate(args): + if isinstance(obj, torch.nn.Module): + model_count += 1 + elif isinstance(obj, torch.optim.Optimizer): + optimizer_count += 1 + + # This needs to be written as such, so that passing other objects other than models/optimizers doesn't raise an error + if (model_count < 1 and optimizer_count > 0) or (model_count > 0 and optimizer_count < 1): + raise ValueError( + "When using FSDP2, a model and optimizer must be passed together to `Accelerator.prepare()`" + " as the optimizer needs to have its parameters modified after the model is converted." + ) + if model_count > 1: + raise ValueError("Only one model is supported when using FSDP2") + + # If we're dealing with device placement, this deals with that by... + tpu_should_fix_optimizer = self.device_placement and self.distributed_type == DistributedType.XLA + + if tpu_should_fix_optimizer: + # 1. grabbing old model parameters + old_named_params = self._get_named_parameters(*args, drop_refs=False) + + if self.distributed_type in [DistributedType.MULTI_CPU, DistributedType.MULTI_XPU, DistributedType.NO]: + if ( + is_torch_version("<", "2.7.0") + and (self.device.type == "cpu" or self.device.type == "xpu") + and self.state.use_ipex + ): + logger.warning( + "You are using lower version of PyTorch(< 2.7.0) with ipex acceleration on Intel CPU or XPU, Intel has upstreamed most of the optimizations into stock PyTorch from 2.7.0, we encourage you to install the latest stock PyTorch and enjoy the out-of-experience on Intel CPU/XPU." + ) + args = self._prepare_ipex(*args) + if self.parallelism_config and self.parallelism_config.tp_enabled: + args = self._prepare_tp(*args) + + if self.parallelism_config and self.parallelism_config.cp_enabled: + args = self._prepare_cp(*args) + + if self.fp8_backend == FP8BackendType.TE: + args = self._prepare_te(*args) + elif self.fp8_backend == FP8BackendType.AO: + args = self._prepare_ao(*args) + if self.distributed_type == DistributedType.DEEPSPEED: + result = self._prepare_deepspeed(*args) + elif self.distributed_type == DistributedType.MEGATRON_LM: + result = self._prepare_megatron_lm(*args) + elif self.is_fsdp2: + result = self._prepare_fsdp2(*args) + else: + if self.fp8_backend == FP8BackendType.MSAMP: + args, device_placement = self._prepare_msamp(*args, device_placement=device_placement) + result = tuple( + self._prepare_one(obj, first_pass=True, device_placement=d) for obj, d in zip(args, device_placement) + ) + result = tuple(self._prepare_one(obj, device_placement=d) for obj, d in zip(result, device_placement)) + if tpu_should_fix_optimizer: + # 2. grabbing new model parameters + new_named_params = self._get_named_parameters(*result) + # 3. building a map from the first to the second + mapping = {p: new_named_params[n] for n, p in old_named_params.items()} + # 4. using that map to update the parameters of the optimizer + for obj in result: + if isinstance(obj, torch.optim.Optimizer): + obj._switch_parameters(mapping) + + for item in result: + if any( + item in container + for container in (self._dataloaders, self._models, self._optimizers, self._schedulers) + ): + item._is_accelerate_prepared = True + + return result if len(result) > 1 else result[0] + + def _prepare_tp(self, *args): + # First pass: prepare everything except schedulers (and model, which is prepared separately below) + result = [ + self._prepare_one(obj, first_pass=True) if not isinstance(obj, torch.nn.Module) else obj for obj in args + ] + + # Second pass: prepare schedulers + result = [self._prepare_one(obj) if not isinstance(obj, torch.nn.Module) else obj for obj in result] + + device_mesh = self.torch_device_mesh + + old_named_params = fsdp2_canonicalize_names(self._get_named_parameters(*tuple(result), drop_refs=True)) + + for arg in result: + if not isinstance(arg, torch.nn.Module): + continue + + from torch.distributed.tensor import DTensor, Replicate + from transformers.integrations.tensor_parallel import ReplicateParallel + + model: torch.nn.Module = arg + tp_plan = ReplicateParallel + + for name, param in model.named_parameters(): + if isinstance(param, DTensor): + continue + + dp = DTensor.from_local(param, device_mesh=device_mesh["tp"], placements=[Replicate()]) + param_name, param_type = name.rsplit(".", 1) + module_to_tp = model.get_submodule(param_name) + + tp_plan().prepare_module_tp(module_to_tp, device_mesh["tp"]) + if not isinstance(dp, torch.nn.Parameter): + dp = torch.nn.Parameter(dp, requires_grad=param.requires_grad) + setattr(module_to_tp, param_type, dp) + + new_named_params = fsdp2_canonicalize_names(self._get_named_parameters(*tuple(result), drop_refs=False)) + # Build a map from old to new params + mapping = {p: new_named_params[n] for n, p in old_named_params.items()} + + def _get_tensor_address(p): + if isinstance(p, DTensor): + return p._local_tensor.data_ptr() + return p.data_ptr() + + for obj in result: + if isinstance(obj, torch.optim.Optimizer): + for param_group in obj.param_groups: + # Each param_group originally maps to model parameters (e.g., from model.parameters()). + # After _prepare_tp(), parameter references are replaced with DTensor instances. + # Therefore, we remap the parameter references to their new DTensor addresses + # so that the optimizer can correctly update the model parameters. + param_group["params"] = [mapping[_get_tensor_address(p)] for p in param_group["params"]] + + return args + + def _prepare_cp(self, *args): + if self.parallelism_config.sp_backend == "deepspeed": + # deepspeed handles cp in a different way, configured in _prepare_deepspeed + return args + + from torch.distributed.tensor.experimental import context_parallel + from torch.distributed.tensor.experimental._attention import set_rotate_method + + cp_comm_strategy = self.parallelism_config.cp_handler.cp_comm_strategy + set_rotate_method(cp_comm_strategy) + + self._cp_context = functools.partial(context_parallel, mesh=self.torch_device_mesh["cp"]) + + for arg in args: + if isinstance(arg, torch.nn.Module): + _attach_context_parallel_hooks(arg) + + return args + + def _prepare_fsdp2(self, *args): + # First pass: prepare everything except schedulers (and model, which is prepared separately below) + result = [ + self._prepare_one(obj, first_pass=True) if not isinstance(obj, torch.nn.Module) else obj for obj in args + ] + + # Second pass: prepare schedulers + result = [self._prepare_one(obj) if not isinstance(obj, torch.nn.Module) else obj for obj in result] + + # Prepare the model + model_index, model = None, None + for i, obj in enumerate(result): + if isinstance(obj, torch.nn.Module): + model_index, model = i, obj + + # Invariant: if we have a model, we also have an optimizer (checked in `prepare`) + if model_index is None: + return tuple(result) + + # Needs to be done first, to make sure AC + fully_shard will work as expected + self.state.fsdp_plugin.set_auto_wrap_policy(model) + + # Apply AC if needed + if self.state.fsdp_plugin.activation_checkpointing: + model = fsdp2_apply_ac(self, model) + + # Apply compile if needed, has to be *after* applying AC + # Copied from: `accelerator.prepare_model` ~ L1804 + if self.state.dynamo_plugin.backend != DynamoBackend.NO and not is_compiled_module(model): + if self.state.dynamo_plugin.use_regional_compilation: + model = compile_regions(model, **self.state.dynamo_plugin.to_kwargs()) + else: + model = torch.compile(model, **self.state.dynamo_plugin.to_kwargs()) + + # Get old params and canonicalize - we canonicalize to have the mapping easy + old_named_params = fsdp2_canonicalize_names(self._get_named_parameters(*tuple(result), drop_refs=True)) + + # Swap the optimizer parameters with empty, so `fully_shard` after will not allocate too much memory + from torch.distributed.tensor import DTensor + + for obj in result: + if isinstance(obj, torch.optim.Optimizer): + for param_group in obj.param_groups: + for i, p in enumerate(param_group["params"]): + # We drop a reference to the original param here, so that _move_states_to_device triggers a reallocation + # We reassign the data_ptr to the original param, so that we preserve the mapping to the new ones + param_group["params"][i] = torch.empty(1, dtype=p.dtype, device=p.device) + param_group["params"][i].data_ptr = ( + p._local_tensor.data_ptr() if isinstance(p, DTensor) else p.data_ptr() + ) + + self._models.append(model) + + # Prepare everything FSDP2 related for the model (except AC) + model = fsdp2_prepare_model(self, model) + + # Remove the old model from the list + if len(self._models) > 1 and (self._models[-2] is self._models[-1]): + del self._models[-2] + + # Replace the old model with the new one (shouldn't be needed as everything should be in place) + result[model_index] = model + + # Get new params and canonicalize + new_named_params = fsdp2_canonicalize_names(self._get_named_parameters(*result)) + # Build a map from old to new params + mapping = {p: new_named_params[n] for n, p in old_named_params.items()} + # Update the optimizer parameters + for obj in result: + if isinstance(obj, torch.optim.Optimizer): + fsdp2_switch_optimizer_parameters(obj, mapping) + + return result + + def prepare_model( + self, model: torch.nn.Module, device_placement: bool | None = None, evaluation_mode: bool = False + ): + """ + Prepares a PyTorch model for training in any distributed setup. It is recommended to use + [`Accelerator.prepare`] instead. + + Args: + model (`torch.nn.Module`): + A PyTorch model to prepare. You don't need to prepare a model if it is used only for inference without + any kind of mixed precision + device_placement (`bool`, *optional*): + Whether or not to place the model on the proper device. Will default to `self.device_placement`. + evaluation_mode (`bool`, *optional*, defaults to `False`): + Whether or not to set the model for evaluation only, by just applying mixed precision and + `torch.compile` (if configured in the `Accelerator` object). + + Example: + + ```python + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator() + >>> # Assume a model is defined + >>> model = accelerator.prepare_model(model) + ``` + """ + if device_placement is None: + device_placement = self.device_placement and self.distributed_type != DistributedType.FSDP + + self._models.append(model) + + # TODO: Look at enabling native TP training directly with a proper config + if ( + self.verify_device_map(model) + and self.distributed_type != DistributedType.NO + and os.environ.get("ACCELERATE_BYPASS_DEVICE_MAP", "false") != "true" + ): + raise ValueError( + "You can't train a model that has been loaded with `device_map='auto'` in any distributed mode." + " Please rerun your script specifying `--num_processes=1` or by launching with `python {{myscript.py}}`." + ) + + if self.native_amp: + model._original_forward = model.forward + autocast_context = get_mixed_precision_context_manager(self.native_amp, self.autocast_handler) + # NOTE: MS-AMP adds `__func__` already to `model.forward`, so we should always use `model.forward` + if self.fp8_backend == FP8BackendType.MSAMP or not hasattr(model.forward, "__func__"): + model_forward_func = model.forward + model.forward = convert_outputs_to_fp32(autocast_context(model_forward_func)) + else: + model_forward_func = model.forward.__func__ + new_forward = autocast_context(model_forward_func) + model.forward = MethodType(new_forward, model) + model.forward = MethodType(convert_outputs_to_fp32(model.forward.__func__), model) + + # We prepare TE after, allowing for bf16 autocast to happen first + if self.fp8_backend == FP8BackendType.TE and not self.delayed_fp8_autocast: + model = apply_fp8_autowrap(model, self.te_recipe_handler or self.fp8_recipe_handler) + + if (getattr(model, "is_loaded_in_8bit", False) or getattr(model, "is_loaded_in_4bit", False)) and getattr( + model, "hf_device_map", False + ): + model_devices = set(model.hf_device_map.values()) + if len(model_devices) > 1 and self.distributed_type != DistributedType.NO: + raise ValueError( + "You can't train a model that has been loaded in 8-bit or 4-bit precision on multiple devices in any distributed mode." + " In order to use 8-bit or 4-bit models that have been loaded across multiple GPUs the solution is to use Naive Pipeline Parallelism." + " Therefore you should not specify that you are under any distributed regime in your accelerate config." + ) + elif len(model_devices) == 1: + current_device = list(model_devices)[0] + if isinstance(current_device, torch.device): + current_device_index = current_device.index + elif isinstance(current_device, str): + current_device_index = torch.device(current_device).index + else: + current_device_index = current_device + + if self.device.type == "cpu" and is_bitsandbytes_multi_backend_available(): + # bnb with multi-backend supports CPU which don't need to check index. + pass + elif torch.device(current_device_index) != self.device: + # if on the first device (GPU 0) we don't care + if (self.device.index is not None) or (current_device_index != 0): + raise ValueError( + "You can't train a model that has been loaded in 8-bit or 4-bit precision on a different device than the one " + "you're training on. Make sure you loaded the model on the correct device using for example `device_map={'':torch.cuda.current_device()}` or `device_map={'':torch.xpu.current_device()}`" + ) + if ( + ("cpu" in model_devices and not is_bitsandbytes_multi_backend_available()) + or ("cpu" in model_devices and is_xpu_available()) + or "disk" in model_devices + ): + raise ValueError( + "You can't train a model that has been loaded in 8-bit or 4-bit precision with CPU or disk offload. " + "If you want train the 8-bit or 4-bit model in CPU, please install bitsandbytes with multi-backend, see https://huggingface.co/docs/bitsandbytes/main/en/installation#multi-backend" + ) + elif device_placement and not self.verify_device_map(model): + model = model.to(self.device) + if not evaluation_mode: + if self.multi_device and not (self.parallelism_config and self.parallelism_config.tp_enabled): + if model_has_dtensor(model): + raise ValueError( + "Your model contains `DTensor` parameters, which is incompatible with DDP. Maybe you loaded your model with `device_map='auto'`? Specify `device_map='cuda'` or 'cpu' instead." + ) + if any(p.requires_grad for p in model.parameters()): + kwargs = self.ddp_handler.to_kwargs() if self.ddp_handler is not None else {} + # TODO: Look at enabling native TP training directly with a proper config + if os.environ.get("ACCELERATE_BYPASS_DEVICE_MAP", "false") != "true": + if self.device.type == "hpu": + device_ids, output_device = [self.device.index], self.device.index + else: + device_ids, output_device = [self.local_process_index], self.local_process_index + else: + device_ids, output_device = None, None + model = torch.nn.parallel.DistributedDataParallel( + model, device_ids=device_ids, output_device=output_device, **kwargs + ) + if self.ddp_handler is not None: + self.ddp_handler.register_comm_hook(model) + elif self.parallelism_config and self.parallelism_config.tp_enabled: + if not hasattr(model, "tp_size"): + raise NotImplementedError( + "Model should undergo tensor parallel before passing it to accelerate." + "You can use .from_pretrained(..., tp_plan='auto') if the model supports" + ) + if model.tp_size != self.parallelism_config.tp_size: + raise ValueError( + f"tp_size in the plugin {self.parallelism_config.tp_size} should be same as model's tp size {model.tp_size}" + ) + elif self.is_fsdp2: + raise ValueError( + "FSDP2 preparation should be done via `accelerate.prepare()`, as it requires a model and an optimizer." + ) + + elif self.distributed_type == DistributedType.FSDP: + # We need to fix the optimizer *before* sharding the model + from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP + + # Check if the model is already a FSDP model due to `Manual Wrapping` and if so, + # don't wrap it again + # In case the model is already compiled using PyTorch 2.0 and the wrapped model in it + # is a FSDP model, don't wrap it again + is_type_fsdp = isinstance(model, FSDP) or ( + is_compiled_module(model) and isinstance(model._orig_mod, FSDP) + ) + + if not is_type_fsdp: + self.state.fsdp_plugin.set_auto_wrap_policy(model) + fsdp_plugin = self.state.fsdp_plugin + + # need to ensure that params are re-tied after running + # param_init_fn + fsdp_plugin.param_init_fn = ensure_weights_retied( + fsdp_plugin.param_init_fn, + model, + self.device, + ) + + kwargs = { + # We fallback to reshard_after_forward if sharding_strategy is not set. + # We prerfer sharding_strategy to not break the behavior of the existing code. + # Deprecation warning has already been issued in `utils.dataclasses.py` + "sharding_strategy": fsdp_plugin.sharding_strategy or fsdp_plugin.reshard_after_forward, + "cpu_offload": fsdp_plugin.cpu_offload, + "auto_wrap_policy": fsdp_plugin.auto_wrap_policy, + "mixed_precision": fsdp_plugin.mixed_precision_policy, + "sync_module_states": fsdp_plugin.sync_module_states, + "backward_prefetch": fsdp_plugin.backward_prefetch, + "forward_prefetch": fsdp_plugin.forward_prefetch, + "use_orig_params": fsdp_plugin.use_orig_params, + "param_init_fn": fsdp_plugin.param_init_fn, + "ignored_modules": fsdp_plugin.ignored_modules, + "limit_all_gathers": fsdp_plugin.limit_all_gathers, + "device_id": self.device, + } + + if isinstance(kwargs["ignored_modules"], str): + reg = re.compile(kwargs["ignored_modules"]) + ignored = [] + for name, module in model.named_modules(): + if reg.fullmatch(name): + # ensure that the device for these modules is still set correctly + module.to(self.device) + ignored.append(module) + kwargs["ignored_modules"] = ignored + + model = FSDP(model, **kwargs) + if fsdp_plugin.activation_checkpointing: + from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ( + CheckpointImpl, + apply_activation_checkpointing, + checkpoint_wrapper, + ) + + apply_activation_checkpointing( + model, + checkpoint_wrapper_fn=functools.partial( + checkpoint_wrapper, + checkpoint_impl=CheckpointImpl.NO_REENTRANT, + ), + auto_wrap_policy=fsdp_plugin.auto_wrap_policy, + ) + + # In the event the model had been loaded in low precision, but + # mixed precision had also been activated, then we follow DeepSpeed's + # strategy to hold the parameters in full precision. + # - assume that trainer.args.bf16 and trainer.args.fp16 are already checked against + # fsdp_plugin.mixed_precision_policy. + # - NOTE: we do not check the mixed_precision attribute on the FSDP root wrapper. + # * this attribute will always set by init_utils.init_core_state so its always not None. + # * mixed_precision.param_dtype only regards _fwd_bwd_param_dtype + # * if model is loaded in 16bit, and even if mixed_precision.param_dtype is None, + # we still want to upcast the flat_param. + if self.mixed_precision != "no": # if mixed precision is set + upcasted_log = [] + for module in FSDP.fsdp_modules(model): + # Referencing DeepSpeed Zero3 + # - in Init, params are converted to 16bit while partitioning. + # - in accelerator.prepare, deepspeed.initialize is called to: + # * creates the DeepSpeedEngine. + # * since zero_optimization() is True , calls engine._configure_zero_optimizer. + # + # Inside the DeepSpeed Zero3 optimizer configuration, which initializes + # DeepSpeedZeroOptimizer_Stage3, during which: + # * trainable_param_groups are obtained from the attached optimizer + # (already partitioned in 16bit). + # * then _setup_for_real_optimizer -> _create_fp32_partitions + # which performs the fp32 upcasting. + + # To mimic DeepSeepds's casting in FSDP, we look at the (single) FlatParameter held + # within an FSDP wrapper. This FlatParameter will be seen by the optimizer. + # - even though there is a torch.device('meta') guard below, we + # expect _init_utils._init_param_handle_from_module to already + # sync the parameter. + + if not module._has_params: + continue # skip if FSDP module not managing parameters + param = module._flat_param + if ( + param.dtype != torch.float32 + and param.device != torch.device("meta") + and param.requires_grad + ): + # keep log of names_params that was upcasted + # NOTE: resorted to this because warnings.simplefilter("once") is somehow not working + name_param_log = (module.module.__class__.__name__, ", ".join(module._flat_param._fqns)) + if name_param_log not in upcasted_log: + upcasted_log.append(name_param_log) + + # this works because of FSDP's _runtime_utils.lazy_init. + # Have to be careful not to call anything before this that + # triggers lazy_init (e.g., _is_fsdp_root). + param.data = param.data.to(torch.float32) # upcasting + module._handle._orig_param_dtype = torch.float32 # update + + # report the warnings + # some messages can be quite repetitive, especially when reporting about layers that have identical architecture. + if self.is_main_process: + for name_log, param_log in upcasted_log: + warnings.warn( + f"Upcasted low precision parameters in {name_log} because mixed precision turned on in FSDP. " + f"Affects: {param_log}." + ) + + if len(upcasted_log) > 0: + warnings.warn( + "FSDP upcast of low precision parameters may affect the precision of model checkpoints." + ) + + # if the previous and current models are same, delete the previous one + if len(self._models) > 1 and (self._models[-2] is self._models[-1]): + del self._models[-2] + self._models[-1] = model + elif self.distributed_type == DistributedType.MULTI_CPU: + kwargs = self.ddp_handler.to_kwargs() if self.ddp_handler else {} + model = torch.nn.parallel.DistributedDataParallel(model, **kwargs) + if self.ddp_handler is not None: + self.ddp_handler.register_comm_hook(model) + elif self.distributed_type == DistributedType.XLA and self.state.fork_launched: + model = xmp.MpModelWrapper(model).to(self.device) + # Now we can apply the FP8 autocast + if self.fp8_backend == FP8BackendType.TE and self.delayed_fp8_autocast: + model = apply_fp8_autowrap(model, self.te_recipe_handler or self.fp8_recipe_handler) + # torch.compile should be called last and only if the model isn't already compiled + if self.state.dynamo_plugin.backend != DynamoBackend.NO and not is_compiled_module(model): + if self.state.dynamo_plugin.use_regional_compilation: + model = compile_regions(model, **self.state.dynamo_plugin.to_kwargs()) + else: + model = torch.compile(model, **self.state.dynamo_plugin.to_kwargs()) + return model + + def _prepare_ao(self, *args): + if not is_torchao_available(): + raise ImportError( + "`torchao` was not found on your system or is too old of a version. Please ensure that `torchao >= 0.6.1` is installed" + ) + + if self.is_fsdp2: + models = [x for x in args if isinstance(x, torch.nn.Module)] + optimizers = [x for x in args if isinstance(x, torch.optim.Optimizer)] + for arg in args: + if isinstance(arg, torch.nn.Module): + convert_model_to_fp8_ao( + arg, + config=self.ao_recipe_handler.config, + module_filter_func=self.ao_recipe_handler.module_filter_func, + ) + + # Invariant: with FSDP2, optimizer is always passed to `prepare()` together with model + # We only precompute scales if float8 all gather is enabled, possibly can add a flag for this later + if self.is_fsdp2 and len(optimizers) > 0 and self.ao_recipe_handler.config.enable_fsdp_float8_all_gather: + from torchao.float8 import precompute_float8_dynamic_scale_for_fsdp + + optimizers[0].register_step_post_hook( + lambda *args, **kwargs: precompute_float8_dynamic_scale_for_fsdp(models[0]) + ) + + return args + + def _prepare_te(self, *args): + if not is_transformer_engine_available(): + raise ImportError( + "`transformer_engine` was not found on your system. Please ensure that `transformer_engine` is installed" + ) + model, optimizer = None, None + num_models, num_optimizers = 0, 0 + result = [obj for obj in args] + for obj in result: + if isinstance(obj, torch.nn.Module): + model = obj + num_models += 1 + elif isinstance(obj, (torch.optim.Optimizer)): + optimizer = obj + num_optimizers += 1 + if optimizer is None and model is None: + return result + elif optimizer is None or model is None: + raise ValueError( + "You must pass a model and an optimizer together to `accelerate.prepare()` when using TransformerEngine." + ) + elif num_models > 1 or num_optimizers > 1: + raise ValueError( + f"You can't use multiple models ({num_models}) or optimizers {num_optimizers} with TransformerEngine." + ) + old_named_params = self._get_named_parameters(model) + with torch.no_grad(): + convert_model(model) + new_named_params = self._get_named_parameters(model) + mapping = {p: new_named_params[n] for n, p in old_named_params.items()} + # We need to switch the optimizer params to the new params *after* the model is wrapped in FSDP + for param_group in optimizer.param_groups: + param_group["params"] = [mapping[p] for p in param_group["params"]] + + return result + + def _prepare_deepspeed(self, *args): + import deepspeed + + ds_initialize = deepspeed.initialize + if self.fp8_backend == FP8BackendType.MSAMP: + # MS-AMP requires DeepSpeed patches + from msamp import deepspeed as msamp_deepspeed + + ds_initialize = msamp_deepspeed.initialize + + deepspeed_plugin = self.deepspeed_plugin + + is_dataloader_present = any(isinstance(obj, torch.utils.data.DataLoader) for obj in args) + tp_size = deepspeed_plugin.deepspeed_config.get("tensor_parallel", {}).get("autotp_size", 0) + + sp_backend = self.parallelism_config.sp_backend if self.parallelism_config else None + sp_size = self.parallelism_config.sp_size if self.parallelism_config else 1 + sp_handler = self.parallelism_config.sp_handler if self.parallelism_config else None + + if tp_size > 1: + if not compare_versions("deepspeed", ">=", "0.16.4"): + raise ImportError( + "Deepspeed TP requires deepspeed >= 0.16.4, Please update DeepSpeed via `pip install deepspeed -U`." + ) + if not is_torch_version(">=", "2.2.0"): + raise ImportError( + "Tried to use TP, but `torch.distributed.device_mesh` requires PyTorch >= 2.2.0. Please upgrade your PyTorch version" + ) + from torch.distributed.device_mesh import init_device_mesh + + mesh_dim_name = "tp" + self.state.ds_device_mesh = init_device_mesh(self.device.type, (tp_size,), mesh_dim_names=(mesh_dim_name,)) + + result = [ + self._prepare_one(obj, first_pass=True) if isinstance(obj, torch.utils.data.DataLoader) else obj + for obj in args + ] + + if deepspeed_plugin.is_auto("train_micro_batch_size_per_gpu"): + if is_dataloader_present: + batch_sizes = [obj.batch_size for obj in args if hasattr(obj, "batch_size")] + if any(bs is None for bs in batch_sizes): + raise ValueError( + "At least one of the dataloaders passed to `accelerate.prepare()` has `None` as batch size. " + "Please set an integer value in `train_micro_batch_size_per_gpu` in the deepspeed config file " + "or assign integer value to `AcceleratorState().deepspeed_plugin.deepspeed_config['train_micro_batch_size_per_gpu']`." + ) + if self.split_batches: + batch_sizes = [batch_size // self.num_processes for batch_size in batch_sizes] + + batch_size_per_device = min(batch_sizes) if deepspeed_plugin.is_train_batch_min else max(batch_sizes) + if len(batch_sizes) > 1: + logger.info( + "Since you passed both train and evaluation dataloader, `is_train_batch_min` (here " + f"{deepspeed_plugin.is_train_batch_min} will decide the `train_batch_size` ({batch_size_per_device})." + ) + else: + raise ValueError( + "When using DeepSpeed, `accelerate.prepare()` requires you to pass at least one of training or evaluation dataloaders " + "with `batch_size` attribute returning an integer value " + "or alternatively set an integer value in `train_micro_batch_size_per_gpu` in the deepspeed config file " + "or assign integer value to `AcceleratorState().deepspeed_plugin.deepspeed_config['train_micro_batch_size_per_gpu']`." + ) + else: + batch_size_per_device = deepspeed_plugin.get_value("train_micro_batch_size_per_gpu") + + # handle `gradient_accumulation_steps` when the value is `auto` + deepspeed_plugin.fill_match( + "gradient_accumulation_steps", + must_match=False, + gradient_accumulation_steps=self.gradient_accumulation_steps, + ) + + deepspeed_gradient_accumulation_steps = deepspeed_plugin.get_value("gradient_accumulation_steps") + # update gradient_accumulation_steps if there is a mismatch + if deepspeed_gradient_accumulation_steps != self.gradient_accumulation_steps: + logger.warning( + f"Gradient accumulation steps mismatch: GradientAccumulationPlugin has {self.gradient_accumulation_steps}, " + f"DeepSpeed config has {deepspeed_gradient_accumulation_steps}. Using DeepSpeed's value." + ) + self.gradient_accumulation_steps = deepspeed_gradient_accumulation_steps + + config_kwargs = { + "gradient_clipping": 1.0, + "zero_optimization.stage3_gather_16bit_weights_on_model_save": False, + } + # This block is skipped when preparing just a model and DL is absent from current call's args + if batch_size_per_device is not None: + config_kwargs["train_micro_batch_size_per_gpu"] = batch_size_per_device + config_kwargs["train_batch_size"] = ( + batch_size_per_device + * deepspeed_plugin.get_value("gradient_accumulation_steps") + * self.num_processes + // sp_size + ) + + model = None + optimizer = None + scheduler = None + for obj in result: + if isinstance(obj, torch.nn.Module): + model = obj + elif isinstance(obj, (torch.optim.Optimizer, DummyOptim)): + optimizer = obj + elif (isinstance(obj, (LRScheduler, DummyScheduler))) or ( + type(obj).__name__ in deepspeed.runtime.lr_schedules.VALID_LR_SCHEDULES + ): + scheduler = obj + + if optimizer is not None: + if "optimizer" in deepspeed_plugin.deepspeed_config and not isinstance(optimizer, (DummyOptim)): + raise ValueError( + "You cannot specify an optimizer in the config file and in the code at the same time. " + "Please remove the optimizer from the config file or " + "create `accelerate.utils.DummyOptim` in the code." + ) + elif "optimizer" not in deepspeed_plugin.deepspeed_config and isinstance(optimizer, (DummyOptim)): + raise ValueError( + "You cannot create a `DummyOptim` without specifying an optimizer in the config file." + ) + + if isinstance(optimizer, (torch.optim.Optimizer)): + deepspeed_plugin.deepspeed_config["zero_allow_untested_optimizer"] = True + + if scheduler is not None: + if "scheduler" in deepspeed_plugin.deepspeed_config and not isinstance(scheduler, (DummyScheduler)): + raise ValueError( + "You cannot specify a scheduler in the config file and in the code at the same time. " + "Please remove the scheduler from the config file or " + "create `accelerate.utils.DummyScheduler` in the code." + ) + elif ( + "scheduler" not in deepspeed_plugin.deepspeed_config + and isinstance(scheduler, (DummyScheduler)) + and scheduler.lr_scheduler_callable is None + ): + raise ValueError( + "Either specify a scheduler in the config file or " + "pass in the `lr_scheduler_callable` parameter when using `accelerate.utils.DummyScheduler`." + ) + + if optimizer is not None and scheduler is not None: + if isinstance(optimizer, (DummyOptim)) and not isinstance(scheduler, (DummyScheduler)): + raise ValueError( + "You can only specify `accelerate.utils.DummyScheduler` in the code when using " + "`accelerate.utils.DummyOptim`." + ) + + if model is not None: + # If we are using FP8, we need to apply the autowrap now + if self.fp8_backend == FP8BackendType.TE: + model = apply_fp8_autowrap(model, self.fp8_recipe_handler) + # if the model is an MOE, set the appropriate MOE layers as leaf Z3 modules + deepspeed_plugin.set_moe_leaf_modules(model) + # deal with config keys that use `auto` value and rely on model's hidden_size + hidden_size_based_keys = [ + "zero_optimization.reduce_bucket_size", + "zero_optimization.stage3_prefetch_bucket_size", + "zero_optimization.stage3_param_persistence_threshold", + ] + hidden_size_auto_keys = [x for x in hidden_size_based_keys if deepspeed_plugin.is_auto(x)] + if len(hidden_size_auto_keys) > 0: + reasoning = ( + "therefore it's not possible to automatically fill out the following `auto` entries " + + f"in the DeepSpeed config file: {hidden_size_auto_keys}. You can fix that by replacing " + + "`auto` values for these keys with an integer value of your choice." + ) + if not hasattr(model, "config"): + raise ValueError("Can't find `model.config` entry, " + reasoning) + + if hasattr(model.config, "hidden_size"): + hidden_size = model.config.hidden_size + elif hasattr(model.config, "hidden_sizes"): + # if there are many hidden sizes pick the largest one + hidden_size = max(model.config.hidden_sizes) + else: + raise ValueError( + "Can find neither `model.config.hidden_size` nor `model.config.hidden_sizes`, " + reasoning + ) + + config_kwargs.update( + { + "zero_optimization.reduce_bucket_size": hidden_size * hidden_size, + "zero_optimization.stage3_prefetch_bucket_size": int(0.9 * hidden_size * hidden_size), + "zero_optimization.stage3_param_persistence_threshold": 10 * hidden_size, + } + ) + + if isinstance(optimizer, (DummyOptim)): + config_kwargs.update( + {"optimizer.params.lr": optimizer.lr, "optimizer.params.weight_decay": optimizer.weight_decay} + ) + if isinstance(scheduler, (DummyScheduler)) and scheduler.lr_scheduler_callable is None: + max_lr = ( + getattr(scheduler.optimizer, "lr", None) + if getattr(scheduler.optimizer, "defaults", None) is None + else scheduler.optimizer.defaults["lr"] + ) + config_kwargs.update( + { + "scheduler.params.warmup_min_lr": 0, + "scheduler.params.warmup_max_lr": max_lr, + "scheduler.params.warmup_num_steps": scheduler.warmup_num_steps, + } + ) + if scheduler.total_num_steps is not None: + config_kwargs["scheduler.params.total_num_steps"] = ( + math.ceil(scheduler.total_num_steps / self.num_processes) + if not self.split_batches + else scheduler.total_num_steps + ) + + deepspeed_plugin.deepspeed_config_process(must_match=False, **config_kwargs) + self.deepspeed_config = deepspeed_plugin.deepspeed_config + + # note: batch_size derivation is all over the map, especiall in HF Trainer, so try to fix it at the last moment if needed + pc = self.parallelism_config + if pc is not None and pc.sp_backend == "deepspeed" and pc.sp_size > 1: + self.deepspeed_config["train_batch_size"] = ( + self.deepspeed_config["train_micro_batch_size_per_gpu"] + * self.deepspeed_config["gradient_accumulation_steps"] + * pc.data_parallel_size + ) + + kwargs = dict(model=model, config_params=self.deepspeed_config) + if optimizer is not None: + if isinstance(optimizer, (DummyOptim)): + kwargs["model_parameters"] = optimizer.params + if isinstance(scheduler, (DummyScheduler)) and scheduler.lr_scheduler_callable is not None: + kwargs["lr_scheduler"] = scheduler.lr_scheduler_callable + else: + if self.deepspeed_config["zero_optimization"].get("offload_optimizer", {}).get( + "device", "none" + ) != "none" and self.deepspeed_config.get("zero_force_ds_cpu_optimizer", True): + if self.device.type == "hpu" and os.environ.get("PT_HPU_LAZY_MODE", "1") == "1": + raise ValueError( + "You can't use an Offload Optimizer with HPU in Lazy Mode. " + "Please set the environment variable `PT_HPU_LAZY_MODE` to `0`." + ) + + optimizer = map_pytorch_optim_to_deepspeed(optimizer) + kwargs["optimizer"] = optimizer + if scheduler is not None: + if type(scheduler).__name__ in deepspeed.runtime.lr_schedules.VALID_LR_SCHEDULES: + kwargs["lr_scheduler"] = scheduler + + if self.device.type == "hpu": + # This env variable is initialized here to make sure it is set to "true" + # It should be done by the launcher but it does not work for multi-node runs + os.environ["DEEPSPEED_USE_HPU"] = "true" + + mpu = None + if sp_size > 1: + if sp_backend != "deepspeed": + raise ValueError( + f"In order to use the configured {sp_size=} with DeepSpeed, you need to configure sp_backend='deepspeed', yet you configured it to be {sp_backend=}." + ) + + ver_min_required = "0.18.2" + if not compare_versions("deepspeed", ">=", ver_min_required): + raise ImportError( + f"Deepspeed ALST/Ulysses requires deepspeed>={ver_min_required}. Please update DeepSpeed via `pip install deepspeed -U`." + ) + + from deepspeed.runtime.sequence_parallel.ulysses_sp import ( + UlyssesSPAttentionHF, + UlyssesSPDataLoaderAdapter, + ) + + if not hasattr(model, "config"): + raise ValueError( + "UlyssesSPAttentionHF currently works with HF Transformers and expects the model object to have a config attribute but this model doesn't have one." + ) + + mpu = UlyssesSPAttentionHF.register_with_transformers( + model_name_or_path=model, + sequence_parallel_size=sp_size, + seq_length=sp_handler.sp_seq_length, + seq_length_is_variable=sp_handler.sp_seq_length_is_variable, + core_attn_implementation=sp_handler.sp_attn_implementation, + micro_batch_size=batch_size_per_device, + ) + kwargs["mpu"] = mpu + + for i in range(len(result)): + if isinstance(result[i], torch.utils.data.DataLoader): + if sp_size > 1: + # note that in case dataloader was prepared apart from model (for the external accelerator.prepare call) you'd need to call deepspeed_ulysses_dl_adapter after prepare(model) (see HF Trainer as the use-case) + sp_group = mpu.get_sequence_parallel_group() + sp_world_size = mpu.get_sequence_parallel_world_size() + sp_rank = mpu.get_sequence_parallel_rank() + result[i] = UlyssesSPDataLoaderAdapter( + result[i], + sp_rank=sp_rank, + sp_group=sp_group, + sp_world_size=sp_world_size, + device=self.device, # model.device, + ) + + engine, optimizer, _, lr_scheduler = ds_initialize(**kwargs) + + if compare_versions("deepspeed", ">=", "0.14.4") and self.state.dynamo_plugin.backend != DynamoBackend.NO: + compile_kwargs = self.state.dynamo_plugin.to_kwargs() + if self.state.dynamo_plugin.use_regional_compilation: + compile_regions_deepspeed(engine.module, **compile_kwargs) + else: + engine.compile(backend=compile_kwargs.pop("backend"), compile_kwargs=compile_kwargs) + if optimizer is not None: + optimizer = DeepSpeedOptimizerWrapper(optimizer) + if scheduler is not None: + if lr_scheduler is None: + scheduler = AcceleratedScheduler( + scheduler, + optimizer, + step_with_optimizer=self.step_scheduler_with_optimizer, + split_batches=self.split_batches, + ) + else: + scheduler = DeepSpeedSchedulerWrapper(lr_scheduler, optimizer) + + for i in range(len(result)): + if isinstance(result[i], torch.nn.Module): + result[i] = engine + elif isinstance(result[i], (torch.optim.Optimizer, DummyOptim)): + result[i] = optimizer + elif (isinstance(result[i], (LRScheduler, DummyScheduler))) or ( + type(result[i]).__name__ in deepspeed.runtime.lr_schedules.VALID_LR_SCHEDULES + ): + result[i] = scheduler + + # pointing for deepspeed_engine_wrapped.backward() + if self.deepspeed_engine_wrapped is None: + self.deepspeed_engine_wrapped = DeepSpeedEngineWrapper(engine) + else: + logger.warning( + "A wrapped DeepSpeed engine reference is currently tied for this `Accelerator()` instance. " + "If you want to call `accelerator.backward()` referencing a new model/engine, " + "please create a separate `Accelerator()` instance and call `accelerator.prepare()` on it." + ) + self._models.append(engine) + if optimizer is not None: + self._optimizers.append(optimizer) + if scheduler is not None: + self._schedulers.append(scheduler) + return tuple(result) + + def deepspeed_ulysses_dl_adapter(self, dl, model): + """this is normally called as part of `prepare` but when dataloader was prepared apart from model (for the external accelerator.prepare call) this additional call needs to be made after prepare(model) (see HF Trainer as the use-case)""" + sp_size = self.parallelism_config.sp_size if self.parallelism_config else 1 + if sp_size == 1: + return dl + from deepspeed.runtime.sequence_parallel.ulysses_sp import UlyssesSPDataLoaderAdapter + from deepspeed.utils import groups + + sp_group = groups._get_sequence_parallel_group() + sp_world_size = groups._get_sequence_parallel_world_size() + sp_rank = groups._get_sequence_parallel_rank() + dl = UlyssesSPDataLoaderAdapter( + dl, + sp_rank=sp_rank, + sp_group=sp_group, + sp_world_size=sp_world_size, + device=model.device, + ) + return dl + + def _prepare_megatron_lm(self, *args): + megatron_lm_plugin = self.state.megatron_lm_plugin + micro_batch_size = None + if not megatron_lm_plugin.megatron_dataset_flag: + batch_sizes = [obj.batch_size for obj in args if hasattr(obj, "batch_size")] + if len(batch_sizes) == 0: + raise ValueError( + "You must specify a training or evaluation dataloader in `accelerate.prepare()` when using Megatron-LM." + ) + + micro_batch_size = min(batch_sizes) if megatron_lm_plugin.is_train_batch_min else max(batch_sizes) + if len(batch_sizes) > 1: + logger.info( + "Since you passed both train and evaluation dataloader, `is_train_batch_min` (here " + f"{megatron_lm_plugin.is_train_batch_min} will decide the `train_batch_size` ({micro_batch_size})." + ) + else: + for obj in args: + if isinstance(obj, MegatronLMDummyDataLoader): + micro_batch_size = obj.dataset_args["micro_batch_size"] + break + if micro_batch_size is not None: + dp_degree = self.num_processes // (megatron_lm_plugin.tp_degree * megatron_lm_plugin.pp_degree) + megatron_lm_plugin.set_training_args(micro_batch_size, dp_degree) + else: + raise ValueError( + "When you do not pass the dataloader parameter, the `data_parallel_size`, " + "`micro_batch_size`, and `global_batch_size` megatron parameters will not be updated." + ) + model = None + optimizer = None + scheduler = None + batch_data = None + for obj in args: + if isinstance(obj, torch.utils.data.DataLoader) and batch_data is None: + batch_data = next(iter(obj)) + elif isinstance(obj, torch.nn.Module): + model = obj + elif isinstance(obj, (torch.optim.Optimizer)): + optimizer = obj + elif isinstance(obj, (LRScheduler, MegatronLMDummyScheduler)): + scheduler = obj + + if model is not None: + megatron_lm_plugin.set_network_size_args(model, batch_data) + if optimizer is not None: + megatron_lm_plugin.set_optimizer_type(optimizer) + if scheduler is not None: + if not isinstance(scheduler, MegatronLMDummyScheduler): + raise ValueError( + "You can't use a custom scheduler with Megatron-LM. Please use the `accelerate.utils.MegatronLMDummyScheduler` instead." + ) + megatron_lm_plugin.set_scheduler_args(scheduler) + + # initialize megatron-lm + megatron_lm_initialize(self, args_defaults=megatron_lm_plugin.megatron_lm_default_args) + + (model, optimizer, scheduler) = megatron_lm_prepare_model_optimizer_scheduler(self) + self.wait_for_everyone() + + counter = 0 + result = [] + for obj in args: + if isinstance(obj, torch.utils.data.DataLoader): + result.append(megatron_lm_prepare_data_loader(self, obj)) + counter += 1 + elif isinstance(obj, MegatronLMDummyDataLoader): + if counter == 0: + obj.set_megatron_data_args() + dataloaders = megatron_lm_prepare_data_loader(self, obj) + result.append(dataloaders[counter]) + counter += 1 + else: + result.append(obj) + + if model is not None: + model = MegatronEngine(self, model, optimizer, scheduler) + if optimizer is not None: + optimizer = MegatronLMOptimizerWrapper(optimizer) + if scheduler is not None: + scheduler = MegatronLMSchedulerWrapper(scheduler, optimizer) + + for i in range(len(result)): + if isinstance(result[i], torch.nn.Module): + result[i] = model + elif isinstance(result[i], torch.optim.Optimizer): + result[i] = optimizer + elif isinstance(result[i], MegatronLMDummyScheduler): + result[i] = scheduler + + if model is not None: + self._models.append(model) + if len(self._models) > 1: + raise AssertionError( + "You can't use same `Accelerator()` instance with multiple models when using Megatron-LM" + ) + if optimizer is not None: + self._optimizers.append(optimizer) + if scheduler is not None: + self._schedulers.append(scheduler) + + return tuple(result) + + def _prepare_ipex(self, *args): + """ + Prepares model and optimizer for training with IPEX on CPU/XPU. This covers 3 cases, IPEX compiled with CPU + only support, IPEX compiled with XPU support and training with XPU pytorch backend available in stock pytorch + starting from version 2.4. + """ + + # ipex.optimize() is available only for IPEX, both IPEX-CPU and IPEX-XPU + if is_ipex_available(): + import intel_extension_for_pytorch as ipex + else: + raise ImportError( + "IPEX is not installed or IPEX's version does not match current PyTorch version. Please refer" + " to https://github.com/intel/intel-extension-for-pytorch." + ) + + models = [] + optimizers = [] + result = [obj for obj in args] + for i, obj in enumerate(result): + if isinstance(obj, torch.nn.Module): + model = obj + model.train() + models.append((i, model)) + elif isinstance(obj, (torch.optim.Optimizer)): + optimizers.append((i, obj)) + + # Impossible to determine what to do if multiple models and/or optimizers are provided + if len(optimizers) > 1 or (len(models) > 1 and len(optimizers) == 1): + raise ValueError( + "Prepare with IPEX expects either 1+ models and no optimizer OR a single model-optimizer pair." + ) + + # Nothing to do + if len(models) == 0 and len(optimizers) == 0: + return result + + dtype = torch.bfloat16 if self.state.mixed_precision == "bf16" else None + # Multiple models and no optimizer (inference) are provided + if len(models) > 0 and len(optimizers) == 0: + for i, model in models: + if self.device.type == "xpu" and next(model.parameters()).device.type == "cpu": + model = model.to(self.device) + model, _ = ipex.optimize(model, optimizer=None, dtype=dtype, inplace=True, level="O1") + # Replace in result + result[i] = model + + # A single model-optimizer pair (training) is provided + if len(models) == 1 and len(optimizers) == 1: + i_model, model = models[0] + i_optimizer, optimizer = optimizers[0] + if self.device.type == "xpu" and next(model.parameters()).device.type == "cpu": + model = model.to(self.device) + model, optimizer = ipex.optimize(model, optimizer=optimizer, dtype=dtype, inplace=True, level="O1") + # Replace in result + result[i_model] = model + result[i_optimizer] = optimizer + + return tuple(result) + + def _prepare_device_mesh(self): + """ + Prepare the device mesh for distributed training. The dataloader will determine how to load data based on the + device mesh. + """ + if self.distributed_type == DistributedType.DEEPSPEED and hasattr(self.state, "ds_device_mesh"): + return self.state.ds_device_mesh + else: + return self.torch_device_mesh + + def _prepare_msamp(self, *args, device_placement): + if not is_msamp_available(): + raise ImportError( + "MS-AMP was not found on your system. Please ensure that MS-AMP is available " + " or choose `'te'` as the backend for FP8 mixed precision training." + ) + # We've already checked for FSDP + MS-AMP during `__init__` + import msamp + + model, optimizer = None, None + optimizer_index = None + num_models, num_optimizers = 0, 0 + result = [obj for obj in args] + for i, obj in enumerate(result): + if isinstance(obj, torch.nn.Module): + model = obj + num_models += 1 + elif isinstance(obj, (torch.optim.Optimizer)): + optimizer = obj + optimizer_index = i + num_optimizers += 1 + # DataLoader/Scheduler case + if optimizer is None and model is None: + return result, device_placement + elif optimizer is None or model is None: + raise ValueError( + "You must pass a model and an optimizer together to `accelerate.prepare()` when using MS-AMP." + ) + elif num_models > 1 or num_optimizers > 1: + raise ValueError( + f"You can't use multiple models ({num_models}) or optimizers {num_optimizers} with MS-AMP." + ) + else: + # DEPRECATE @ 2.0 + if self.fp8_recipe_handler is not None: + opt_level = self.fp8_recipe_handler.opt_level + else: + opt_level = self.msamp_recipe_handler.opt_level + model, optimizer = msamp.initialize(model, optimizer, opt_level=opt_level) + for i in range(len(result)): + if isinstance(result[i], torch.nn.Module): + result[i] = model + elif isinstance(result[i], (torch.optim.Optimizer)): + result[i] = optimizer + if optimizer_index is not None: + # NOTE: MS-AMP moves the optimizer, but *not* the model to the right device + device_placement[optimizer_index] = False + return tuple(result), device_placement + + def prepare_data_loader( + self, data_loader: torch.utils.data.DataLoader, device_placement=None, slice_fn_for_dispatch=None + ): + """ + Prepares a PyTorch DataLoader for training in any distributed setup. It is recommended to use + [`Accelerator.prepare`] instead. + + Args: + data_loader (`torch.utils.data.DataLoader`): + A vanilla PyTorch DataLoader to prepare + device_placement (`bool`, *optional*): + Whether or not to place the batches on the proper device in the prepared dataloader. Will default to + `self.device_placement`. + slice_fn_for_dispatch (`Callable`, *optional*`): + If passed, this function will be used to slice tensors across `num_processes`. Will default to + [`~utils.slice_tensors`]. This argument is used only when `dispatch_batches` is set to `True` and will + be ignored otherwise. + + Example: + + ```python + >>> import torch + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator() + >>> data_loader = torch.utils.data.DataLoader(...) + >>> data_loader = accelerator.prepare_data_loader(data_loader, device_placement=True) + ``` + """ + # Ensure we can't double wrap a DataLoader due to `find_batch_size` + if getattr(data_loader, "_is_accelerate_prepared", False): + if data_loader not in self._dataloaders: + self._dataloaders.append(data_loader) + return data_loader + if device_placement is None: + device_placement = self.device_placement if self.distributed_type != DistributedType.XLA else False + + device_mesh = self._prepare_device_mesh() + + prepared_data_loader = prepare_data_loader( + data_loader, + self.device, + num_processes=self.num_processes, + process_index=self.process_index, + split_batches=self.split_batches, + put_on_device=device_placement, + rng_types=self.rng_types.copy(), + dispatch_batches=self.dispatch_batches, + even_batches=self.even_batches, + slice_fn_for_dispatch=slice_fn_for_dispatch, + use_seedable_sampler=self.use_seedable_sampler, + data_seed=self.dataloader_config.data_seed, + non_blocking=self.non_blocking, + use_stateful_dataloader=self.use_stateful_dataloader, + torch_device_mesh=device_mesh, + ) + self._dataloaders.append(prepared_data_loader) + return prepared_data_loader + + def prepare_optimizer(self, optimizer: torch.optim.Optimizer, device_placement=None): + """ + Prepares a PyTorch Optimizer for training in any distributed setup. It is recommended to use + [`Accelerator.prepare`] instead. + + Args: + optimizer (`torch.optim.Optimizer`): + A vanilla PyTorch optimizer to prepare + device_placement (`bool`, *optional*): + Whether or not to place the optimizer on the proper device. Will default to `self.device_placement`. + + Example: + + ```python + >>> import torch + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator() + >>> optimizer = torch.optim.Adam(...) + >>> optimizer = accelerator.prepare_optimizer(optimizer, device_placement=True) + ``` + """ + if is_lomo_available(): + # We need to import locally to avoid circular imports since lomo imports stuff from + # transformers & accelerate + from lomo_optim import AdaLomo, Lomo + + # Support multiple optimizers: https://github.com/huggingface/accelerate/pull/2695#discussion_r1589164607 + self.has_lomo_optimizer |= isinstance(optimizer, (Lomo, AdaLomo)) + + # Ensure we can't double wrap an optimizer due to `find_batch_size` + if getattr(optimizer, "_is_accelerate_prepared", False): + if optimizer not in self._optimizers: + self._optimizers.append(optimizer) + return optimizer + if device_placement is None: + device_placement = self.device_placement + # NOTE: Special case with MS-AMP we do *not* pass in the scaler explicitly to the `AcceleratedOptimizer`, + # Their optimizer handles it for us. + scaler = None if self.fp8_backend == FP8BackendType.MSAMP else self.scaler + optimizer = AcceleratedOptimizer(optimizer, device_placement=device_placement, scaler=scaler) + self._optimizers.append(optimizer) + return optimizer + + def prepare_scheduler(self, scheduler: LRScheduler): + """ + Prepares a PyTorch Scheduler for training in any distributed setup. It is recommended to use + [`Accelerator.prepare`] instead. + + Args: + scheduler (`torch.optim.lr_scheduler.LRScheduler`): + A vanilla PyTorch scheduler to prepare + + Example: + + ```python + >>> import torch + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator() + >>> optimizer = torch.optim.Adam(...) + >>> scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, ...) + >>> scheduler = accelerator.prepare_scheduler(scheduler) + ``` + """ + # Ensure we can't double wrap a scheduler due to `find_batch_size` + if getattr(scheduler, "_is_accelerate_prepared", False): + if scheduler not in self._schedulers: + self._schedulers.append(scheduler) + return scheduler + # We try to find the optimizer associated with `scheduler`, the default is the full list. + optimizer = self._optimizers + for opt in self._optimizers: + if getattr(scheduler, "optimizer", None) == opt.optimizer: + optimizer = opt + break + scheduler = AcceleratedScheduler( + scheduler, + optimizer, + step_with_optimizer=self.step_scheduler_with_optimizer, + split_batches=self.split_batches, + ) + self._schedulers.append(scheduler) + return scheduler + + def backward(self, loss, **kwargs): + """ + Scales the gradients in accordance to the `GradientAccumulationPlugin` and calls the correct `backward()` based + on the configuration. + + Should be used in lieu of `loss.backward()`. + + Example: + + ```python + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator(gradient_accumulation_steps=2) + >>> outputs = model(inputs) + >>> loss = loss_fn(outputs, labels) + >>> accelerator.backward(loss) + ``` + """ + learning_rate = kwargs.get("learning_rate") + + if self.distributed_type != DistributedType.DEEPSPEED: + # deepspeed handles loss scaling by gradient_accumulation_steps in its `backward` + loss = loss / self.gradient_accumulation_steps + if self.distributed_type == DistributedType.DEEPSPEED: + self.deepspeed_engine_wrapped.backward(loss, sync_gradients=self.sync_gradients, **kwargs) + elif self.distributed_type == DistributedType.MEGATRON_LM: + return + elif self.scaler is not None: + self.scaler.scale(loss).backward(**kwargs) + elif learning_rate is not None and self.has_lomo_optimizer: + self.lomo_backward(loss, learning_rate) + else: + loss.backward(**kwargs) + + def set_trigger(self): + """ + Sets the internal trigger tensor to 1 on the current process. A latter check should follow using this which + will check across all processes. + + Note: + Does not require `wait_for_everyone()` + + Example: + + ```python + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator() + >>> # Assume later in the training script + >>> # `should_do_breakpoint` is a custom function to monitor when to break, + >>> # e.g. when the loss is NaN + >>> if should_do_breakpoint(loss): + ... accelerator.set_trigger() + >>> # Assume later in the training script + >>> if accelerator.check_breakpoint(): + ... break + ``` + """ + self.flag_tensor = torch.tensor(1, device=self.device) + + def check_trigger(self): + """ + Checks if the internal trigger tensor has been set to 1 in any of the processes. If so, will return `True` and + reset the trigger tensor to 0. + + Note: + Does not require `wait_for_everyone()` + + Example: + + ```python + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator() + >>> # Assume later in the training script + >>> # `should_do_breakpoint` is a custom function to monitor when to break, + >>> # e.g. when the loss is NaN + >>> if should_do_breakpoint(loss): + ... accelerator.set_trigger() + >>> # Assume later in the training script + >>> if accelerator.check_trigger(): + ... break + ``` + """ + # Now that we are outside `__init__`, we can initialize it if it is `None` on device + if self.flag_tensor is None: + self.flag_tensor = torch.tensor(0, device=self.device) + flag_tensor = self.reduce(self.flag_tensor) + if flag_tensor.item() >= 1: + self.flag_tensor = torch.tensor(0, device=self.device) + return True + return False + + def unscale_gradients(self, optimizer=None): + """ + Unscale the gradients in mixed precision training with AMP. This is a noop in all other settings. + + Likely should be called through [`Accelerator.clip_grad_norm_`] or [`Accelerator.clip_grad_value_`] + + Args: + optimizer (`torch.optim.Optimizer` or `list[torch.optim.Optimizer]`, *optional*): + The optimizer(s) for which to unscale gradients. If not set, will unscale gradients on all optimizers + that were passed to [`~Accelerator.prepare`]. + + Example: + + ```python + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator() + >>> model, optimizer = accelerator.prepare(model, optimizer) + >>> outputs = model(inputs) + >>> loss = loss_fn(outputs, labels) + >>> accelerator.backward(loss) + >>> accelerator.unscale_gradients(optimizer=optimizer) + ``` + """ + if self.native_amp and self.mixed_precision == "fp16": + if optimizer is None: + # TODO: this unscales all optimizers where we should only unscale the one where parameters are. + optimizer = self._optimizers + elif not isinstance(optimizer, (tuple, list)): + optimizer = [optimizer] + for opt in optimizer: + while isinstance(opt, AcceleratedOptimizer): + opt = opt.optimizer + self.scaler.unscale_(opt) + + def clip_grad_norm_(self, parameters, max_norm, norm_type=2): + """ + Should be used in place of `torch.nn.utils.clip_grad_norm_`. + + Returns: + `torch.Tensor`: Total norm of the parameter gradients (viewed as a single vector). + + Example: + + ```python + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator(gradient_accumulation_steps=2) + >>> dataloader, model, optimizer, scheduler = accelerator.prepare(dataloader, model, optimizer, scheduler) + + >>> for input, target in dataloader: + ... optimizer.zero_grad() + ... output = model(input) + ... loss = loss_func(output, target) + ... accelerator.backward(loss) + ... if accelerator.sync_gradients: + ... accelerator.clip_grad_norm_(model.parameters(), max_grad_norm) + ... optimizer.step() + ``` + """ + if self.distributed_type == DistributedType.FSDP: + self.unscale_gradients() + parameters = [p for p in parameters] + for model in self._models: + if parameters == [p for p in model.parameters()]: + if not self.is_fsdp2: + return model.clip_grad_norm_(max_norm, norm_type) + else: + return torch.nn.utils.clip_grad_norm_( + parameters, max_norm, norm_type=norm_type + ) # viz: https://github.com/pytorch/torchtitan/blob/main/docs/fsdp.md + elif self.distributed_type == DistributedType.DEEPSPEED: + # DeepSpeed handles gradient clipping internally, but we can retrieve the gradient norm + if self.deepspeed_engine_wrapped is not None: + return self.deepspeed_engine_wrapped.get_global_grad_norm() + return None + elif self.distributed_type == DistributedType.XLA: + # Reduce gradients first for XLA + for acc_opt in self._optimizers: + if not acc_opt.gradient_state.is_xla_gradients_synced: + opt = acc_opt + while isinstance(opt, AcceleratedOptimizer): + opt = opt.optimizer + gradients = xm._fetch_gradients(opt) + # Use xm.all_reduce to perform an in-place all-reduce. Recursive all-reduce each tensor + # one by one in self.reduce is non-inplace. + xm.all_reduce("sum", gradients, scale=1.0 / self.num_processes) + # Set is_xla_gradients_synced to True to avoid all-reduce twice in the AcceleratedOptimizer step. + acc_opt.gradient_state.is_xla_gradients_synced = True + if os.environ.get("ACCELERATE_USE_FSDP", "false").lower() == "true": + self.unscale_gradients() + parameters = [p for p in parameters] + for model in self._models: + if parameters == [p for p in model.parameters()]: + return model.clip_grad_norm_(max_norm, norm_type) + self.unscale_gradients() + return torch.nn.utils.clip_grad_norm_(parameters, max_norm, norm_type=norm_type) + + def clip_grad_value_(self, parameters, clip_value): + """ + Should be used in place of `torch.nn.utils.clip_grad_value_`. + + Example: + + ```python + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator(gradient_accumulation_steps=2) + >>> dataloader, model, optimizer, scheduler = accelerator.prepare(dataloader, model, optimizer, scheduler) + + >>> for input, target in dataloader: + ... optimizer.zero_grad() + ... output = model(input) + ... loss = loss_func(output, target) + ... accelerator.backward(loss) + ... if accelerator.sync_gradients: + ... accelerator.clip_grad_value_(model.parameters(), clip_value) + ... optimizer.step() + ``` + """ + if self.distributed_type in [DistributedType.DEEPSPEED, DistributedType.FSDP]: + raise Exception("DeepSpeed and FSDP do not support `clip_grad_value_`. Use `clip_grad_norm_` instead.") + self.unscale_gradients() + torch.nn.utils.clip_grad_value_(parameters, clip_value) + + def gather(self, tensor): + """ + Gather the values in *tensor* across all processes and concatenate them on the first dimension. Useful to + regroup the predictions from all processes when doing evaluation. + + Note: + This gather happens in all processes. + + Args: + tensor (`torch.Tensor`, or a nested tuple/list/dictionary of `torch.Tensor`): + The tensors to gather across all processes. + + Returns: + `torch.Tensor`, or a nested tuple/list/dictionary of `torch.Tensor`: The gathered tensor(s). Note that the + first dimension of the result is *num_processes* multiplied by the first dimension of the input tensors. + + Example: + + ```python + >>> # Assuming four processes + >>> import torch + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator() + >>> process_tensor = torch.tensor([accelerator.process_index], device=accelerator.device) + >>> gathered_tensor = accelerator.gather(process_tensor) + >>> gathered_tensor + tensor([0, 1, 2, 3]) + ``` + """ + return gather(tensor) + + def gather_for_metrics(self, input_data, use_gather_object=False): + """ + Gathers `input_data` and potentially drops duplicates in the last batch if on a distributed system. Should be + used for gathering the inputs and targets for metric calculation. + + Args: + input (`torch.Tensor`, `object`, a nested tuple/list/dictionary of `torch.Tensor`, or a nested tuple/list/dictionary of `object`): + The tensors or objects for calculating metrics across all processes + use_gather_object(`bool`): + Whether to forcibly use gather_object instead of gather (which is already done if all objects passed do + not contain tensors). This flag can be useful for gathering tensors with different sizes that we don't + want to pad and concatenate along the first dimension. Using it with GPU tensors is not well supported + and inefficient as it incurs GPU -> CPU transfer since tensors would be pickled. + + Example: + + ```python + >>> # Assuming two processes, with a batch size of 5 on a dataset with 9 samples + >>> import torch + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator() + >>> dataloader = torch.utils.data.DataLoader(range(9), batch_size=5) + >>> dataloader = accelerator.prepare(dataloader) + >>> batch = next(iter(dataloader)) + >>> gathered_items = accelerator.gather_for_metrics(batch) + >>> len(gathered_items) + 9 + ``` + """ + + try: + recursively_apply(lambda x: x, input_data, error_on_other_type=True) + all_tensors = True + except TypeError: + all_tensors = False + + use_gather_object = use_gather_object or not all_tensors + + if use_gather_object: + data = gather_object(input_data) + else: + data = self.gather(input_data) + + try: + if self.gradient_state.end_of_dataloader: + # at the end of a dataloader, `gather_for_metrics` regresses to + # `gather` unless the dataset has a remainder so log. + if self.gradient_state.remainder == -1: + logger.info( + "The used dataset had no length, returning gathered tensors. You should drop the remainder yourself." + ) + return data + elif self.gradient_state.remainder > 0: + # Last batch needs to be truncated on distributed systems as it contains additional samples + def _adjust_samples(tensor): + return tensor[: self.gradient_state.remainder] + + if use_gather_object: + # gather_object put the objects in a list + return _adjust_samples(data) + else: + return recursively_apply(_adjust_samples, data) + else: # remainder is 0 + # no remainder even though at end of dataloader, so nothing to do. + return data + else: + # Not at the end of the dataloader, no need to adjust the tensors + return data + except Exception: + # Dataset had no length or raised an error + return data + + def reduce(self, tensor, reduction="sum", scale=1.0): + """ + Reduce the values in *tensor* across all processes based on *reduction*. + + Note: + All processes get the reduced value. + + Args: + tensor (`torch.Tensor`, or a nested tuple/list/dictionary of `torch.Tensor`): + The tensors to reduce across all processes. + reduction (`str`, *optional*, defaults to "sum"): + A reduction type, can be one of 'sum', 'mean', or 'none'. If 'none', will not perform any operation. + scale (`float`, *optional*, defaults to 1.0): + A default scaling value to be applied after the reduce, only valid on XLA. + + Returns: + `torch.Tensor`, or a nested tuple/list/dictionary of `torch.Tensor`: + The reduced tensor(s). + + Example: + + ```python + >>> # Assuming two processes + >>> import torch + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator() + >>> process_tensor = torch.arange(accelerator.num_processes) + 1 + (2 * accelerator.process_index) + >>> process_tensor = process_tensor.to(accelerator.device) + >>> reduced_tensor = accelerator.reduce(process_tensor, reduction="sum") + >>> reduced_tensor + tensor([4, 6]) + ``` + """ + return reduce(tensor, reduction, scale) + + def pad_across_processes(self, tensor, dim=0, pad_index=0, pad_first=False): + """ + Recursively pad the tensors in a nested list/tuple/dictionary of tensors from all devices to the same size so + they can safely be gathered. + + Args: + tensor (nested list/tuple/dictionary of `torch.Tensor`): + The data to gather. + dim (`int`, *optional*, defaults to 0): + The dimension on which to pad. + pad_index (`int`, *optional*, defaults to 0): + The value with which to pad. + pad_first (`bool`, *optional*, defaults to `False`): + Whether to pad at the beginning or the end. + + Returns: + `torch.Tensor`, or a nested tuple/list/dictionary of `torch.Tensor`: + The padded tensor(s). + + Example: + + ```python + >>> # Assuming two processes, with the first processes having a tensor of size 1 and the second of size 2 + >>> import torch + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator() + >>> process_tensor = torch.arange(accelerator.process_index + 1).to(accelerator.device) + >>> padded_tensor = accelerator.pad_across_processes(process_tensor) + >>> padded_tensor.shape + torch.Size([2]) + ``` + """ + return pad_across_processes(tensor, dim=dim, pad_index=pad_index, pad_first=pad_first) + + def unwrap_model(self, model, keep_fp32_wrapper: bool = True, keep_torch_compile: bool = True): + """ + Unwraps the `model` from the additional layer possible added by [`~Accelerator.prepare`]. Useful before saving + the model. + + Args: + model (`torch.nn.Module`): + The model to unwrap. + keep_fp32_wrapper (`bool`, *optional*, defaults to `True`): + Whether to not remove the mixed precision hook if it was added. + keep_torch_compile (`bool`, *optional*, defaults to `True`): + Whether to not unwrap compiled model if compiled. + Returns: + `torch.nn.Module`: The unwrapped model. + + Example: + + ```python + >>> # Assuming two GPU processes + >>> from torch.nn.parallel import DistributedDataParallel + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator() + >>> model = accelerator.prepare(MyModel()) + >>> print(model.__class__.__name__) + DistributedDataParallel + + >>> model = accelerator.unwrap_model(model) + >>> print(model.__class__.__name__) + MyModel + ``` + """ + return extract_model_from_parallel(model, keep_fp32_wrapper, keep_torch_compile) + + def wait_for_everyone(self): + """ + Will stop the execution of the current process until every other process has reached that point (so this does + nothing when the script is only run in one process). Useful to do before saving a model. + + Example: + + ```python + >>> # Assuming two GPU processes + >>> import time + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator() + >>> if accelerator.is_main_process: + ... time.sleep(2) + >>> else: + ... print("I'm waiting for the main process to finish its sleep...") + >>> accelerator.wait_for_everyone() + >>> # Should print on every process at the same time + >>> print("Everyone is here") + ``` + """ + wait_for_everyone() + + @on_main_process + def init_trackers(self, project_name: str, config: dict | None = None, init_kwargs: dict | None = {}): + """ + Initializes a run for all trackers stored in `self.log_with`, potentially with starting configurations + + Args: + project_name (`str`): + The name of the project. All trackers will save their data based on this + config (`dict`, *optional*): + Optional starting configuration to be logged. + init_kwargs (`dict`, *optional*): + A nested dictionary of kwargs to be passed to a specific tracker's `__init__` function. Should be + formatted like so: + ```python + {"wandb": {"tags": ["tag_a", "tag_b"]}} + ``` + + Example: + + ```python + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator(log_with="tensorboard") + >>> accelerator.init_trackers( + ... project_name="my_project", + ... config={"learning_rate": 0.001, "batch_size": 32}, + ... init_kwargs={"tensorboard": {"flush_secs": 60}}, + ... ) + ``` + """ + for tracker in self.log_with: + if issubclass(type(tracker), GeneralTracker): + # Custom trackers are already initialized + self.trackers.append(tracker) + else: + tracker_init = LOGGER_TYPE_TO_CLASS[str(tracker)] + if tracker_init.requires_logging_directory: + # We can skip this check since it was done in `__init__` + self.trackers.append( + tracker_init(project_name, self.logging_dir, **init_kwargs.get(str(tracker), {})) + ) + else: + self.trackers.append(tracker_init(project_name, **init_kwargs.get(str(tracker), {}))) + + for tracker in self.trackers: + tracker.start() + + if config is not None: + for tracker in self.trackers: + tracker.store_init_configuration(config) + + def get_tracker(self, name: str, unwrap: bool = False): + """ + Returns a `tracker` from `self.trackers` based on `name` on the main process only. + + Args: + name (`str`): + The name of a tracker, corresponding to the `.name` property. + unwrap (`bool`): + Whether to return the internal tracking mechanism or to return the wrapped tracker instead + (recommended). + + Returns: + `GeneralTracker`: The tracker corresponding to `name` if it exists. + + Example: + + ```python + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator(log_with="tensorboard") + >>> accelerator.init_trackers("my_project") + >>> tensorboard_tracker = accelerator.get_tracker("tensorboard") + ``` + """ + if len(self.trackers) > 0: + for tracker in self.trackers: + if tracker.name == name: + return tracker.tracker if unwrap else tracker + raise ValueError(f"{name} is not an available tracker stored inside the `Accelerator`.") + # Handle tracker only made on main process + return GeneralTracker(_blank=True) + + @on_main_process + def log(self, values: dict, step: int | None = None, log_kwargs: dict | None = {}): + """ + Logs `values` to all stored trackers in `self.trackers` on the main process only. + + Args: + values (`dict`): + Values should be a dictionary-like object containing only types `int`, `float`, or `str`. + step (`int`, *optional*): + The run step. If included, the log will be affiliated with this step. + log_kwargs (`dict`, *optional*): + A nested dictionary of kwargs to be passed to a specific tracker's `log` function. Should be formatted + like so: + ```python + {"wandb": {"tags": ["tag_a", "tag_b"]}} + ``` + + Example: + + ```python + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator(log_with="tensorboard") + >>> accelerator.init_trackers("my_project") + >>> accelerator.log({"loss": 0.5, "accuracy": 0.9}) + ``` + """ + for tracker in self.trackers: + tracker.log(values, step=step, **log_kwargs.get(tracker.name, {})) + + def end_training(self): + """ + Runs any special end training behaviors, such as stopping trackers on the main process only or destoying + process group. Should always be called at the end of your script if using experiment tracking. + + Example: + + ```python + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator(log_with="tensorboard") + >>> accelerator.init_trackers("my_project") + >>> # Do training + >>> accelerator.end_training() + ``` + """ + for tracker in self.trackers: + tracker.finish() + + self.state.destroy_process_group() + + def save(self, obj, f, safe_serialization=False): + """ + Save the object passed to disk once per machine. Use in place of `torch.save`. + + Args: + obj (`object`): The object to save. + f (`str` or `os.PathLike`): Where to save the content of `obj`. + safe_serialization (`bool`, *optional*, defaults to `False`): Whether to save `obj` using `safetensors` + + Note: + If `save_on_each_node` was passed in as a `ProjectConfiguration`, will save the object once per node, + rather than only once on the main node. + + Example: + + ```python + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator() + >>> arr = [0, 1, 2, 3] + >>> accelerator.save(arr, "array.pkl") + ``` + """ + save( + obj, + f, + save_on_each_node=self.project_configuration.save_on_each_node, + safe_serialization=safe_serialization, + ) + + def save_model( + self, + model: torch.nn.Module, + save_directory: Union[str, os.PathLike], + max_shard_size: Union[int, str] = "10GB", + safe_serialization: bool = True, + ): + """ + Save a model so that it can be re-loaded using load_checkpoint_in_model + + Arguments: + model: (`torch.nn.Module`): + Model to be saved. The model can be wrapped or unwrapped. + save_directory (`str` or `os.PathLike`): + Directory to which to save. Will be created if it doesn't exist. + max_shard_size (`int` or `str`, *optional*, defaults to `"10GB"`): + The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size + lower than this size. If expressed as a string, needs to be digits followed by a unit (like `"5MB"`). + + + + If a single weight of the model is bigger than `max_shard_size`, it will be in its own checkpoint shard + which will be bigger than `max_shard_size`. + + + + safe_serialization (`bool`, *optional*, defaults to `True`): + Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`). + + Example: + + ```python + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator() + >>> model = ... + >>> accelerator.save_model(model, save_directory) + ``` + """ + + if os.path.isfile(save_directory): + logger.error(f"Provided path ({save_directory}) should be a directory, not a file") + return + + # get the state_dict of the model + if any(has_offloaded_params(module) for module in model.modules()): + state_dict = get_state_dict_offloaded_model(model) + else: + if any(param.device == torch.device("meta") for param in model.parameters()): + raise RuntimeError("You can't save the model since some parameters are on the meta device.") + state_dict = self.get_state_dict(model) + + # Case: DeepSpeed zero3 gets gathered and `state_dict` is empty + if state_dict is None: + return + os.makedirs(save_directory, exist_ok=True) + + if safe_serialization: + state_dict = clean_state_dict_for_safetensors(state_dict) + weights_name = SAFE_WEIGHTS_NAME if safe_serialization else WEIGHTS_NAME + filename_pattern = SAFE_WEIGHTS_PATTERN_NAME if safe_serialization else WEIGHTS_PATTERN_NAME + + state_dict_split = split_torch_state_dict_into_shards( + state_dict, filename_pattern=filename_pattern, max_shard_size=max_shard_size + ) + + # Clean the folder from a previous save + for filename in os.listdir(save_directory): + full_filename = os.path.join(save_directory, filename) + # If we have a shard file that is not going to be replaced, we delete it, but only from the main process + # in distributed settings to avoid race conditions. + weights_no_suffix = weights_name.replace(".bin", "") + + # make sure that file to be deleted matches format of sharded file, e.g. pytorch_model-00001-of-00005 + filename_no_suffix = filename.replace(".bin", "") + reg = re.compile(r"(.*?)-\d{5}-of-\d{5}") + + if ( + filename.startswith(weights_no_suffix) + and os.path.isfile(full_filename) + and filename not in state_dict_split.filename_to_tensors.keys() + and reg.fullmatch(filename_no_suffix) is not None + and PartialState().is_main_process + ): + os.remove(full_filename) + + # Save the model + for filename, tensors in state_dict_split.filename_to_tensors.items(): + shard = {tensor: state_dict[tensor] for tensor in tensors} + self.save(shard, os.path.join(save_directory, filename), safe_serialization=safe_serialization) + + # Save index if sharded + if state_dict_split.is_sharded: + index = { + "metadata": state_dict_split.metadata, + "weight_map": state_dict_split.tensor_to_filename, + } + save_index_file = SAFE_WEIGHTS_INDEX_NAME if safe_serialization else WEIGHTS_INDEX_NAME + save_index_file = os.path.join(save_directory, save_index_file) + with open(save_index_file, "w", encoding="utf-8") as f: + content = json.dumps(index, indent=2, sort_keys=True) + "\n" + f.write(content) + logger.info( + f"The model is bigger than the maximum size per checkpoint ({max_shard_size}) and is going to be " + f"split in {len(state_dict_split.filename_to_tensors)} checkpoint shards. You can find where each parameters has been saved in the " + f"index located at {save_index_file}." + ) + else: + path_to_weights = os.path.join(save_directory, WEIGHTS_NAME) + logger.info(f"Model weights saved in {path_to_weights}") + + def register_save_state_pre_hook(self, hook: Callable[..., None]) -> hooks.RemovableHandle: + """ + Registers a pre hook to be run before `save_checkpoint` is called in [`Accelerator.save_state`]. + + Args: + hook (`Callable`): + A function to be called in [`Accelerator.save_state`] before `save_checkpoint`. + + The hook should have the following signature: + + `hook(models: list[torch.nn.Module], weights: list[dict[str, torch.Tensor]], input_dir: str) -> None` + + The `models` argument are the models as saved in the accelerator state under `accelerator._models`, `weights` + argument are the state dicts of the `models`, and the `input_dir` argument is the `input_dir` argument passed + to [`Accelerator.load_state`]. + + + + Should only be used in conjunction with [`Accelerator.register_load_state_pre_hook`]. Can be useful to save + configurations in addition to model weights. Can also be used to overwrite model saving with a customized + method. In this case, make sure to remove already loaded weights from the weights list. + + + + Returns: + `torch.utils.hooks.RemovableHandle`: a handle that can be used to remove the added hook by calling + `handle.remove()` + """ + handle = hooks.RemovableHandle(self._save_model_state_pre_hook) + self._save_model_state_pre_hook[handle.id] = hook + return handle + + def save_state(self, output_dir: str | None = None, safe_serialization: bool = True, **save_model_func_kwargs): + """ + Saves the current states of the model, optimizer, scaler, RNG generators, and registered objects to a folder. + + If a `ProjectConfiguration` was passed to the `Accelerator` object with `automatic_checkpoint_naming` enabled + then checkpoints will be saved to `self.project_dir/checkpoints`. If the number of current saves is greater + than `total_limit` then the oldest save is deleted. Each checkpoint is saved in separate folders named + `checkpoint_`. + + Otherwise they are just saved to `output_dir`. + + + + Should only be used when wanting to save a checkpoint during training and restoring the state in the same + environment. + + + + Args: + output_dir (`str` or `os.PathLike`): + The name of the folder to save all relevant weights and states. + safe_serialization (`bool`, *optional*, defaults to `True`): + Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`). + save_model_func_kwargs (`dict`, *optional*): + Additional keyword arguments for saving model which can be passed to the underlying save function, such + as optional arguments for DeepSpeed's `save_checkpoint` function. + + Example: + + ```python + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator() + >>> model, optimizer, lr_scheduler = ... + >>> model, optimizer, lr_scheduler = accelerator.prepare(model, optimizer, lr_scheduler) + >>> accelerator.save_state(output_dir="my_checkpoint") + ``` + """ + if self.project_configuration.automatic_checkpoint_naming: + output_dir = os.path.join(self.project_dir, "checkpoints") + os.makedirs(output_dir, exist_ok=True) + if self.project_configuration.automatic_checkpoint_naming: + folders = [os.path.join(output_dir, folder) for folder in os.listdir(output_dir)] + if ( + self.project_configuration.total_limit is not None + and (len(folders) + 1 > self.project_configuration.total_limit) + and self.is_main_process + ): + + def _inner(folder): + return list(map(int, re.findall(r"[\/]?([0-9]+)(?=[^\/]*$)", folder)))[0] + + folders.sort(key=_inner) + logger.warning( + f"Deleting {len(folders) + 1 - self.project_configuration.total_limit} checkpoints to make room for new checkpoint." + ) + for folder in folders[: len(folders) + 1 - self.project_configuration.total_limit]: + shutil.rmtree(folder) + output_dir = os.path.join(output_dir, f"checkpoint_{self.save_iteration}") + if os.path.exists(output_dir): + raise ValueError( + f"Checkpoint directory {output_dir} ({self.save_iteration}) already exists. Please manually override `self.save_iteration` with what iteration to start with." + ) + self.wait_for_everyone() + os.makedirs(output_dir, exist_ok=True) + logger.info(f"Saving current state to {output_dir}") + + if self.distributed_type == DistributedType.XLA: + # Finish running the previous step before checkpointing + xm.mark_step() + + # Save the models taking care of FSDP and DeepSpeed nuances + weights = [] + for i, model in enumerate(self._models): + if self.distributed_type == DistributedType.FSDP: + logger.info("Saving FSDP model") + save_fsdp_model(self.state.fsdp_plugin, self, model, output_dir, i) + logger.info(f"FSDP Model saved to output dir {output_dir}") + elif self.distributed_type == DistributedType.DEEPSPEED: + logger.info("Saving DeepSpeed Model and Optimizer") + ckpt_id = f"{MODEL_NAME}" if i == 0 else f"{MODEL_NAME}_{i}" + model.save_checkpoint(output_dir, ckpt_id, **save_model_func_kwargs) + logger.info(f"DeepSpeed Model and Optimizer saved to output dir {os.path.join(output_dir, ckpt_id)}") + elif self.distributed_type == DistributedType.MEGATRON_LM: + logger.info("Saving Megatron-LM Model, Optimizer and Scheduler") + model.save_checkpoint(output_dir) + logger.info(f"Megatron-LM Model , Optimizer and Scheduler saved to output dir {output_dir}") + else: + weights.append(self.get_state_dict(model, unwrap=False)) + + # Save the optimizers taking care of FSDP and DeepSpeed nuances + optimizers = [] + if self.distributed_type == DistributedType.FSDP: + for i, opt in enumerate(self._optimizers): + logger.info("Saving FSDP Optimizer") + save_fsdp_optimizer(self.state.fsdp_plugin, self, opt, self._models[i], output_dir, i) + logger.info(f"FSDP Optimizer saved to output dir {output_dir}") + elif self.distributed_type not in [DistributedType.DEEPSPEED, DistributedType.MEGATRON_LM]: + optimizers = self._optimizers + + # Save the lr schedulers taking care of DeepSpeed nuances + schedulers = [] + if self.distributed_type == DistributedType.DEEPSPEED: + for i, scheduler in enumerate(self._schedulers): + if isinstance(scheduler, DeepSpeedSchedulerWrapper): + continue + schedulers.append(scheduler) + elif self.distributed_type not in [DistributedType.MEGATRON_LM]: + schedulers = self._schedulers + + # Save the samplers of the dataloaders + dataloaders = self._dataloaders + + # Call model loading hooks that might have been registered with + # accelerator.register_model_state_hook + for hook in self._save_model_state_pre_hook.values(): + hook(self._models, weights, output_dir) + + save_location = save_accelerator_state( + output_dir, + weights, + optimizers, + schedulers, + dataloaders, + self.state.process_index, + self.step, + self.scaler, + save_on_each_node=self.project_configuration.save_on_each_node, + safe_serialization=safe_serialization, + ) + for i, obj in enumerate(self._custom_objects): + save_custom_state(obj, output_dir, i, save_on_each_node=self.project_configuration.save_on_each_node) + self.project_configuration.iteration += 1 + return save_location + + def register_load_state_pre_hook(self, hook: Callable[..., None]) -> hooks.RemovableHandle: + """ + Registers a pre hook to be run before [`load_checkpoint`] is called in [`Accelerator.load_state`]. + + Args: + hook (`Callable`): + A function to be called in [`Accelerator.load_state`] before `load_checkpoint`. + + The hook should have the following signature: + + `hook(models: list[torch.nn.Module], input_dir: str) -> None` + + The `models` argument are the models as saved in the accelerator state under `accelerator._models`, and the + `input_dir` argument is the `input_dir` argument passed to [`Accelerator.load_state`]. + + + + Should only be used in conjunction with [`Accelerator.register_save_state_pre_hook`]. Can be useful to load + configurations in addition to model weights. Can also be used to overwrite model loading with a customized + method. In this case, make sure to remove already loaded models from the models list. + + + + Returns: + `torch.utils.hooks.RemovableHandle`: a handle that can be used to remove the added hook by calling + `handle.remove()` + """ + handle = hooks.RemovableHandle(self._load_model_state_pre_hook) + self._load_model_state_pre_hook[handle.id] = hook + return handle + + def load_state(self, input_dir: str | None = None, load_kwargs: dict | None = None, **load_model_func_kwargs): + """ + Loads the current states of the model, optimizer, scaler, RNG generators, and registered objects. + + + + Should only be used in conjunction with [`Accelerator.save_state`]. If a file is not registered for + checkpointing, it will not be loaded if stored in the directory. + + + + Args: + input_dir (`str` or `os.PathLike`): + The name of the folder all relevant weights and states were saved in. Can be `None` if + `automatic_checkpoint_naming` is used, and will pick up from the latest checkpoint. + load_kwargs (`dict`, *optional*): + Additional keyword arguments for the underlying `load` function, such as optional arguments for + state_dict and optimizer on. + load_model_func_kwargs (`dict`, *optional*): + Additional keyword arguments for loading model which can be passed to the underlying load function, + such as optional arguments for DeepSpeed's `load_checkpoint` function or a `map_location` to load the + model and optimizer on. + + Example: + + ```python + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator() + >>> model, optimizer, lr_scheduler = ... + >>> model, optimizer, lr_scheduler = accelerator.prepare(model, optimizer, lr_scheduler) + >>> accelerator.load_state("my_checkpoint") + ``` + """ + if input_dir is not None: + # Check if folder exists + input_dir = os.path.expanduser(input_dir) + if not os.path.isdir(input_dir): + raise ValueError(f"Tried to find {input_dir} but folder does not exist") + elif self.project_configuration.automatic_checkpoint_naming: + # Pick up from automatic checkpoint naming + input_dir = os.path.join(self.project_dir, "checkpoints") + folders = [os.path.join(input_dir, folder) for folder in os.listdir(input_dir)] + + def _inner(folder): + return list(map(int, re.findall(r"[\/]?([0-9]+)(?=[^\/]*$)", folder)))[0] + + folders.sort(key=_inner) + input_dir = folders[-1] + else: + raise ValueError("No input_dir provided and automatic checkpoint naming is disabled.") + logger.info(f"Loading states from {input_dir}") + + # Load the models taking care of FSDP and DeepSpeed nuances + models = [] + for i, model in enumerate(self._models): + if self.distributed_type == DistributedType.FSDP: + logger.info("Loading FSDP model") + load_fsdp_model(self.state.fsdp_plugin, self, model, input_dir, i) + logger.info(f"FSDP Model loaded from input dir {input_dir}") + elif self.distributed_type == DistributedType.DEEPSPEED: + logger.info("Loading DeepSpeed Model and Optimizer") + ckpt_id = f"{MODEL_NAME}" if i == 0 else f"{MODEL_NAME}_{i}" + model.load_checkpoint(input_dir, ckpt_id, **load_model_func_kwargs) + logger.info(f"DeepSpeed Model and Optimizer loaded from input dir {os.path.join(input_dir, ckpt_id)}") + elif self.distributed_type == DistributedType.MEGATRON_LM: + logger.info("Loading Megatron-LM Model, Optimizer and Scheduler") + model.load_checkpoint(input_dir) + logger.info(f"Megatron-LM Model , Optimizer and Scheduler loaded from input dir {input_dir}") + else: + models.append(model) + + # We need to load the scaler state before the optimizer for FSDP2 + # (`torch.distributed.checkpoint.set_optimizer_state_dict`) which we use to set the state of the optimizer calls `optimizer.step` on + # a dummy tensor, but since the scaler is not initialized, it will raise an error (the scaler exists but its `_scale` is None) + scaler = None + if self.scaler is not None and self.is_fsdp2: + input_scaler_file = os.path.join(input_dir, SCALER_NAME) + scaler_state = torch.load(input_scaler_file) + self.scaler.load_state_dict(scaler_state) + # We also need to call the `_lazy_init_scale_growth_tracker` to initialize the scaler, as it would else be called + # on the first call to scale + self.scaler._lazy_init_scale_growth_tracker(self.scaler._device) + logger.info("GradScaler state loaded successfully") + else: + scaler = self.scaler + + # Load the optimizers taking care of FSDP and DeepSpeed nuances + optimizers = [] + if self.distributed_type == DistributedType.FSDP: + for i, opt in enumerate(self._optimizers): + logger.info("Loading FSDP Optimizer") + load_fsdp_optimizer(self.state.fsdp_plugin, self, opt, self._models[i], input_dir, i) + logger.info(f"FSDP Optimizer loaded from input dir {input_dir}") + elif self.distributed_type not in [DistributedType.DEEPSPEED, DistributedType.MEGATRON_LM]: + optimizers = self._optimizers + + # Load the lr schedulers taking care of DeepSpeed nuances + schedulers = [] + if self.distributed_type == DistributedType.DEEPSPEED: + for i, scheduler in enumerate(self._schedulers): + if isinstance(scheduler, DeepSpeedSchedulerWrapper): + continue + schedulers.append(scheduler) + elif self.distributed_type not in [DistributedType.MEGATRON_LM]: + schedulers = self._schedulers + + dataloaders = self._dataloaders + + # Call model loading hooks that might have been registered with + # accelerator.register_model_state_hook + for hook in self._load_model_state_pre_hook.values(): + hook(models, input_dir) + + map_location = load_model_func_kwargs.pop("map_location", None) + if map_location is None: + if self.num_processes > 1 and self.multi_device and self.distributed_type != DistributedType.MULTI_XPU: + map_location = "on_device" + else: + map_location = "cpu" + + override_attributes = load_accelerator_state( + input_dir, + models, + optimizers, + schedulers, + dataloaders, + self.state.process_index, + scaler, + map_location, + load_kwargs, + **load_model_func_kwargs, + ) + if "step" in override_attributes: + self.step = override_attributes["step"] + custom_checkpoints = [ + f for f in os.listdir(input_dir) if re.search(r"^custom_checkpoint_\d+\.pkl$", f) is not None + ] + if len(custom_checkpoints) != len(self._custom_objects): + err = ( + f"Number of custom checkpoints in folder {input_dir} does not match the number of registered objects:" + ) + err += f"\n\tFound checkpoints: {len(custom_checkpoints)}" + err += f"\n\tRegistered objects: {len(self._custom_objects)}\n" + err += "Please make sure to only load checkpoints from folders that were created with the same set of registered objects," + err += "or avoid using `custom_checkpoint` in the filename for files in that same directory and load them in manually." + raise RuntimeError(err) + else: + logger.info(f"Loading in {len(custom_checkpoints)} custom states") + for index, obj in enumerate(self._custom_objects): + load_custom_state(obj, input_dir, index) + + def free_memory(self, *objects): + """ + Will release all references to the internal objects stored and call the garbage collector. You should call this + method between two trainings with different models/optimizers. Also will reset `Accelerator.step` to 0. + + Example: + + ```python + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator() + >>> model, optimizer, scheduler = ... + >>> model, optimizer, scheduler = accelerator.prepare(model, optimizer, scheduler) + >>> model, optimizer, scheduler = accelerator.free_memory(model, optimizer, scheduler) + ``` + """ + # Deepspeed needs a bit more prep that should be done first + if hasattr(self, "deepspeed_engine_wrapped"): + if self.deepspeed_engine_wrapped is not None: + self.deepspeed_engine_wrapped.engine.destroy() + self.deepspeed_engine_wrapped = None + objects = release_memory(*objects) + self._schedulers = [] + self._optimizers = [] + self._models = [] + self._dataloaders = [] + self.step = 0 + return objects + + def clear(self, *objects): + """ + Alias for [`Accelerate.free_memory`], releases all references to the internal objects stored and call the + garbage collector. You should call this method between two trainings with different models/optimizers. + + Example: + + ```python + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator() + >>> model, optimizer, scheduler = ... + >>> model, optimizer, scheduler = accelerator.prepare(model, optimizer, scheduler) + >>> model, optimizer, scheduler = accelerator.clear(model, optimizer, scheduler) + ``` + """ + return self.free_memory(*objects) + + def _get_named_parameters(self, *args, drop_refs=False): + named_parameters = {} + accessor_mapping = {} + for obj in args: + if isinstance(obj, torch.nn.Module): + obj = extract_model_from_parallel(obj) + if not drop_refs: + named_parameters.update({n: p for n, p in obj.named_parameters()}) + continue + + # we need this bit as `WeightWithDynamic...` returns 0 when `data_ptr()` is called, + # the underlying pointer is actually hidden in `_tensor` attribute + if self.fp8_backend == FP8BackendType.AO: + from torchao.float8.fsdp_utils import WeightWithDynamicFloat8CastTensor + + accessor_mapping[WeightWithDynamicFloat8CastTensor] = "_tensor" + # we know we're in FSDP2 so DTensor is available + if self.is_fsdp2: + from torch.distributed.tensor import DTensor + + accessor_mapping[DTensor] = "_local_tensor" + + named_parameters.update( + { + n: getattr(p, accessor_mapping[type(p)]).data_ptr() + if type(p) in accessor_mapping + else p.data_ptr() + for n, p in obj.named_parameters() + } + ) + return named_parameters + + def _get_devices(self, *args): + model_device = None + optimizer_device = None + for obj in args: + # Loop through model parameters and stop at the first once we have its device. + if isinstance(obj, torch.nn.Module): + for param in obj.parameters(): + model_device = param.device + break + # Loop through optimizer parameters groups and stop at the first once we have its device. + if isinstance(obj, torch.optim.Optimizer): + for param_group in obj.param_groups: + if len(param_group["params"]) > 0: + optimizer_device = param_group["params"][0].device + break + return (model_device, optimizer_device) + + def get_state_dict(self, model, unwrap=True): + """ + Returns the state dictionary of a model sent through [`Accelerator.prepare`] potentially without full + precision. + + Args: + model (`torch.nn.Module`): + A PyTorch model sent through [`Accelerator.prepare`] + unwrap (`bool`, *optional*, defaults to `True`): + Whether to return the original underlying state_dict of `model` or to return the wrapped state_dict + + Returns: + `dict`: The state dictionary of the model potentially without full precision. + + Example: + + ```python + >>> import torch + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator() + >>> net = torch.nn.Linear(2, 2) + >>> net = accelerator.prepare(net) + >>> state_dict = accelerator.get_state_dict(net) + ``` + """ + + if self.distributed_type == DistributedType.DEEPSPEED: + zero3_sharding = self.deepspeed_config["zero_optimization"]["stage"] == 3 + tp_sharding = self.deepspeed_config.get("tensor_parallel", {}).get("autotp_size", 0) > 1 + if zero3_sharding or tp_sharding: + if model.zero_gather_16bit_weights_on_model_save(): + ver_min_required = "0.16.4" + if tp_sharding and not compare_versions("deepspeed", ">=", ver_min_required): + raise ImportError( + f"Deepspeed TP requires deepspeed>={ver_min_required}. Please update DeepSpeed via `pip install deepspeed -U`." + ) + state_dict = ( + model._consolidated_16bit_state_dict() + if tp_sharding + else model._zero3_consolidated_16bit_state_dict() + ) + else: + raise ValueError( + "Cannot get 16bit model weights because `stage3_gather_16bit_weights_on_model_save` in DeepSpeed config is False. " + "To save the model weights in 16bit, set `stage3_gather_16bit_weights_on_model_save` to True in DeepSpeed config file or " + "set `zero3_save_16bit_model` to True when using `accelerate config`. " + "To save the full checkpoint, run `model.save_checkpoint(save_dir)` and use `zero_to_fp32.py` to recover weights." + ) + else: + from deepspeed.checkpoint.utils import clone_tensors_for_torch_save + + state_dict = clone_tensors_for_torch_save(self.unwrap_model(model).state_dict()) + elif self.is_fsdp2: + from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict + + options = StateDictOptions(full_state_dict=True, broadcast_from_rank0=True, cpu_offload=True) + state_dict = get_model_state_dict(model, options=options) + elif self.distributed_type == DistributedType.FSDP: + from torch.distributed.fsdp import FullStateDictConfig, StateDictType + from torch.distributed.fsdp import FullyShardedDataParallel as FSDP + + full_state_dict_config = FullStateDictConfig(offload_to_cpu=True, rank0_only=True) + with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, full_state_dict_config): + state_dict = model.state_dict() + else: + if unwrap: + model = self.unwrap_model(model) + state_dict = model.state_dict() + + return state_dict + + def register_for_checkpointing(self, *objects): + """ + Makes note of `objects` and will save or load them in during `save_state` or `load_state`. + + These should be utilized when the state is being loaded or saved in the same script. It is not designed to be + used in different scripts. + + + + Every `object` must have a `load_state_dict` and `state_dict` function to be stored. + + + + Example: + + ```python + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator() + >>> # Assume `CustomObject` has a `state_dict` and `load_state_dict` function. + >>> obj = CustomObject() + >>> accelerator.register_for_checkpointing(obj) + >>> accelerator.save_state("checkpoint.pt") + ``` + """ + invalid_objects = [] + for obj in objects: + if not hasattr(obj, "state_dict") or not hasattr(obj, "load_state_dict"): + invalid_objects.append(obj) + if len(invalid_objects) > 0: + err = "All `objects` must include a `state_dict` and `load_state_dict` function to be stored. The following inputs are invalid:" + for index, obj in enumerate(invalid_objects): + err += f"\n\t- Item at index {index}, `{get_pretty_name(obj)}`" + raise ValueError(err) + self._custom_objects.extend(objects) + + @contextmanager + def maybe_context_parallel( + self, + buffers: list[torch.Tensor] | None = None, + buffer_seq_dims: list[int] | None = None, + no_restore_buffers: set[torch.Tensor] | None = None, + ): + """ + A context manager that enables context parallel training. + + Args: + buffers (`list[torch.Tensor]`, `optional`): + Buffers, which are going to be sharded along the sequence dimension. Common examples are inputs, labels + or positional embedding buffers. This context manager will modify these buffers in-place, and after + exiting the context, the buffers will be restored to their original state. To avoid unnecessary + restores, you can use `no_restore_buffers` to specify which buffers don't need to be restored. + buffer_seq_dims (`list[int]`, `optional`): + Sequence dimensions of `buffers`. + no_restore_buffers (`set[torch.Tensor]`, `optional`): + This set must be a subset of `buffers`. Specifies which buffers from `buffers` argument won't be + restored after the context exits. These buffers will be then kept in sharded state. + + + + `context_parallel` is currently supported with FSDP2 and requires `parallelism_config.cp_size` > + 1. If either of these conditions are not met, this context manager will have no effect, though to enable fewer + code changes it will not raise an Exception. + + + + + + This context manager has to be recreated with each training step, as shown in the example below. + + + + Example: + + ```python + >>> for batch in dataloader: + ... with accelerator.maybe_context_parallel( + ... buffers=[batch["input_ids"], batch["attention_mask"]], + ... buffer_seq_dims=[1, 1], + ... no_restore_buffers={batch["input_ids"]}, + ... ): + ... outputs = model(batch) + ... ... + ``` + """ + # We don't need to check FSDP2 as parallelism_config does that for us + # Invariant: in this branch self._cp_context is set, as it was set by `self._prepare_cp` + if ( + self.parallelism_config + and self.parallelism_config.cp_backend == "torch" + and self.parallelism_config.cp_enabled + ): + with self._cp_context( + buffers=buffers, buffer_seq_dims=buffer_seq_dims, no_restore_buffers=no_restore_buffers + ): + yield + else: + logger.warning_once( + "Context parallel training is not enabled. This context manager will have no effect. " + "To enable it, set `parallelism_config.cp_size` > 1 in the `Accelerator` constructor." + ) + yield + + @contextmanager + def autocast(self, autocast_handler: AutocastKwargs = None): + """ + Will apply automatic mixed-precision inside the block inside this context manager, if it is enabled. Nothing + different will happen otherwise. + + A different `autocast_handler` can be passed in to override the one set in the `Accelerator` object. This is + useful in blocks under `autocast` where you want to revert to fp32. + + Example: + + ```python + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator(mixed_precision="fp16") + >>> with accelerator.autocast(): + ... train() + ``` + """ + if autocast_handler is None: + autocast_handler = self.autocast_handler + autocast_context = get_mixed_precision_context_manager(self.native_amp, autocast_handler) + with autocast_context: + yield + + @contextmanager + def profile(self, profile_handler: ProfileKwargs | None = None): + """ + Will profile the code inside the context manager. The profile will be saved to a Chrome Trace file if + `profile_handler.output_trace_dir` is set. + + A different `profile_handler` can be passed in to override the one set in the `Accelerator` object. + + Args: + profile_handler (`ProfileKwargs`, *optional*): + The profile handler to use for this context manager. If not passed, will use the one set in the + `Accelerator` object. + + Example: + + ```python + # Profile with default settings + from accelerate import Accelerator + from accelerate.utils import ProfileKwargs + + accelerator = Accelerator() + with accelerator.profile() as prof: + train() + accelerator.print(prof.key_averages().table()) + + + # Profile with the custom handler + def custom_handler(prof): + print(prof.key_averages().table(sort_by="self_cpu_time_total", row_limit=10)) + + + kwargs = ProfileKwargs(schedule_option=dict(wait=1, warmup=1, active=1), on_trace_ready=custom_handler) + accelerator = Accelerator(kwarg_handler=[kwargs]) + with accelerator.profile() as prof: + for _ in range(10): + train_iteration() + prof.step() + + + # Profile and export to Chrome Trace + kwargs = ProfileKwargs(output_trace_dir="output_trace") + accelerator = Accelerator(kwarg_handler=[kwargs]) + with accelerator.profile(): + train() + ``` + """ + profile_handler = profile_handler or self.profile_handler or ProfileKwargs() + + with profile_handler.build() as profiler: + yield profiler + + if profile_handler.output_trace_dir is None: + return + + os.makedirs(profile_handler.output_trace_dir, exist_ok=True) + profiler.export_chrome_trace( + os.path.join(profile_handler.output_trace_dir, PROFILE_PATTERN_NAME.format(suffix=self.process_index)) + ) + self.wait_for_everyone() + + @property + def optimizer_step_was_skipped(self): + """ + Whether or not the optimizer update was skipped (because of gradient overflow in mixed precision), in which + case the learning rate should not be changed. + """ + for optimizer in self._optimizers: + if optimizer.step_was_skipped: + return True + return False + + def skip_first_batches(self, dataloader, num_batches: int = 0): + """ + Creates a new `torch.utils.data.DataLoader` that will efficiently skip the first `num_batches`. + + Args: + dataloader (`torch.utils.data.DataLoader`): The data loader in which to skip batches. + num_batches (`int`, *optional*, defaults to 0): The number of batches to skip + + Example: + + ```python + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator() + >>> dataloader, model, optimizer, scheduler = accelerator.prepare(dataloader, model, optimizer, scheduler) + >>> skipped_dataloader = accelerator.skip_first_batches(dataloader, num_batches=2) + >>> # for the first epoch only + >>> for input, target in skipped_dataloader: + ... optimizer.zero_grad() + ... output = model(input) + ... loss = loss_func(output, target) + ... accelerator.backward(loss) + ... optimizer.step() + + >>> # subsequent epochs + >>> for input, target in dataloader: + ... optimizer.zero_grad() + ... ... + ``` + """ + return skip_first_batches(dataloader, num_batches=num_batches) + + def __deepcopy__(self, memo): + logger.info("Deep copying the `Accelerator` object, note that this will point to the same original object.") + return self + + def verify_device_map(self, model: torch.nn.Module) -> bool: + """ + Verifies that `model` has not been prepared with big model inference with a device-map resembling `auto`. + """ + # Checks if any of the child modules has the attribute `hf_device_map` and this map has more than one entry. + for m in model.modules(): + if hasattr(m, "hf_device_map") and len(m.hf_device_map) > 1: + return True + + return False + + def lomo_backward(self, loss: torch.Tensor, learning_rate: float) -> None: + """ + Runs backward pass on LOMO optimizers. + """ + if is_lomo_available(): + # We need to import locally to avoid circular imports since lomo imports stuff from + # transformers & accelerate + from lomo_optim import AdaLomo, Lomo + + if learning_rate is None: + raise ValueError("A learning rate must be passed in order to call backward pass with LOMO optimizers.") + + _backward_called = False + + for optimizer in self._optimizers: + if isinstance(optimizer.optimizer, (Lomo, AdaLomo)): + optimizer.optimizer.fused_backward(loss, learning_rate) + _backward_called = True + + if not _backward_called: + raise ValueError( + "Backward pass not properly called on LOMO optimizers. Are you sure you passed a LOMO optimizer in accelerator.prepare()?" + ) + + @property + def fp8_backend(self) -> FP8BackendType: + "Returns the configured backend for training in FP8" + if self.has_fp8_handler: + if self.fp8_recipe_handler is not None: + return FP8BackendType(self.fp8_recipe_handler.backend) + elif self.ao_recipe_handler is not None: + return FP8BackendType.AO + elif self.te_recipe_handler is not None: + return FP8BackendType.TE + elif self.msamp_recipe_handler is not None: + return FP8BackendType.MSAMP + elif self.state.deepspeed_plugin is not None and self.state.deepspeed_plugin.enable_msamp: + return FP8BackendType.MSAMP + + return FP8BackendType(parse_choice_from_env("ACCELERATE_FP8_BACKEND", "NO")) diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/big_modeling.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/big_modeling.py new file mode 100644 index 0000000000000000000000000000000000000000..aa77b09ed554459d28ec1a3a215853e662796a6d --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/big_modeling.py @@ -0,0 +1,789 @@ +# Copyright 2022 The HuggingFace Team. All rights reserved. +# +# 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. + +import logging +import os +import re +from contextlib import contextmanager +from functools import wraps +from typing import Optional, Union + +import torch +import torch.nn as nn + +from .hooks import ( + AlignDevicesHook, + CpuOffload, + LayerwiseCastingHook, + UserCpuOffloadHook, + add_hook_to_module, + attach_align_device_hook, + attach_align_device_hook_on_blocks, +) +from .utils import ( + OffloadedWeightsLoader, + check_cuda_p2p_ib_support, + check_device_map, + extract_submodules_state_dict, + find_tied_parameters, + get_balanced_memory, + infer_auto_device_map, + is_bnb_available, + is_mlu_available, + is_musa_available, + is_npu_available, + is_sdaa_available, + is_xpu_available, + load_checkpoint_in_model, + offload_state_dict, + parse_flag_from_env, + retie_parameters, +) +from .utils.constants import SUPPORTED_PYTORCH_LAYERS_FOR_UPCASTING +from .utils.other import recursive_getattr + + +logger = logging.getLogger(__name__) + + +@contextmanager +def init_empty_weights(include_buffers: Optional[bool] = None): + """ + A context manager under which models are initialized with all parameters on the meta device, therefore creating an + empty model. Useful when just initializing the model would blow the available RAM. + + Args: + include_buffers (`bool`, *optional*): + Whether or not to also put all buffers on the meta device while initializing. + + Example: + + ```python + import torch.nn as nn + from accelerate import init_empty_weights + + # Initialize a model with 100 billions parameters in no time and without using any RAM. + with init_empty_weights(): + tst = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)]) + ``` + + + + Any model created under this context manager has no weights. As such you can't do something like + `model.to(some_device)` with it. To load weights inside your empty model, see [`load_checkpoint_and_dispatch`]. + Make sure to overwrite the default device_map param for [`load_checkpoint_and_dispatch`], otherwise dispatch is not + called. + + + """ + if include_buffers is None: + include_buffers = parse_flag_from_env("ACCELERATE_INIT_INCLUDE_BUFFERS", False) + with init_on_device(torch.device("meta"), include_buffers=include_buffers) as f: + yield f + + +@contextmanager +def init_on_device(device: torch.device, include_buffers: Optional[bool] = None): + """ + A context manager under which models are initialized with all parameters on the specified device. + + Args: + device (`torch.device`): + Device to initialize all parameters on. + include_buffers (`bool`, *optional*): + Whether or not to also put all buffers on the meta device while initializing. + + Example: + + ```python + import torch.nn as nn + from accelerate import init_on_device + + with init_on_device(device=torch.device("cuda")): + tst = nn.Linear(100, 100) # on `cuda` device + ``` + """ + if include_buffers is None: + include_buffers = parse_flag_from_env("ACCELERATE_INIT_INCLUDE_BUFFERS", False) + + if include_buffers: + with device: + yield + return + + old_register_parameter = nn.Module.register_parameter + if include_buffers: + old_register_buffer = nn.Module.register_buffer + + def register_empty_parameter(module, name, param): + old_register_parameter(module, name, param) + if param is not None: + param_cls = type(module._parameters[name]) + kwargs = module._parameters[name].__dict__ + kwargs["requires_grad"] = param.requires_grad + module._parameters[name] = param_cls(module._parameters[name].to(device), **kwargs) + + def register_empty_buffer(module, name, buffer, persistent=True): + old_register_buffer(module, name, buffer, persistent=persistent) + if buffer is not None: + module._buffers[name] = module._buffers[name].to(device) + + # Patch tensor creation + if include_buffers: + tensor_constructors_to_patch = { + torch_function_name: getattr(torch, torch_function_name) + for torch_function_name in ["empty", "zeros", "ones", "full"] + } + else: + tensor_constructors_to_patch = {} + + def patch_tensor_constructor(fn): + def wrapper(*args, **kwargs): + kwargs["device"] = device + return fn(*args, **kwargs) + + return wrapper + + try: + nn.Module.register_parameter = register_empty_parameter + if include_buffers: + nn.Module.register_buffer = register_empty_buffer + for torch_function_name in tensor_constructors_to_patch.keys(): + setattr(torch, torch_function_name, patch_tensor_constructor(getattr(torch, torch_function_name))) + yield + finally: + nn.Module.register_parameter = old_register_parameter + if include_buffers: + nn.Module.register_buffer = old_register_buffer + for torch_function_name, old_torch_function in tensor_constructors_to_patch.items(): + setattr(torch, torch_function_name, old_torch_function) + + +def cpu_offload( + model: nn.Module, + execution_device: Optional[torch.device] = None, + offload_buffers: bool = False, + state_dict: Optional[dict[str, torch.Tensor]] = None, + preload_module_classes: Optional[list[str]] = None, +): + """ + Activates full CPU offload for a model. As a result, all parameters of the model will be offloaded and only one + copy of the state dict of the model will be kept. During the forward pass, parameters will be extracted from that + state dict and put on the execution device passed as they are needed, then offloaded again. + + Args: + model (`torch.nn.Module`): + The model to offload. + execution_device (`torch.device`, *optional*): + The device on which the forward pass of the model will be executed (should be a GPU). Will default to the + model first parameter device. + offload_buffers (`bool`, *optional*, defaults to `False`): + Whether or not to offload the buffers with the model parameters. + state_dict (`Dict[str, torch.Tensor]`, *optional*): + The state dict of the model that will be kept on CPU. + preload_module_classes (`List[str]`, *optional*): + A list of classes whose instances should load all their weights (even in the submodules) at the beginning + of the forward. This should only be used for classes that have submodules which are registered but not + called directly during the forward, for instance if a `dense` linear layer is registered, but at forward, + `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly. + """ + if execution_device is None: + execution_device = next(iter(model.parameters())).device + if state_dict is None: + state_dict = {n: p.to("cpu") for n, p in model.state_dict().items()} + + add_hook_to_module(model, AlignDevicesHook(io_same_device=True), append=True) + attach_align_device_hook( + model, + execution_device=execution_device, + offload=True, + offload_buffers=offload_buffers, + weights_map=state_dict, + preload_module_classes=preload_module_classes, + ) + + return model + + +def cpu_offload_with_hook( + model: torch.nn.Module, + execution_device: Optional[Union[int, str, torch.device]] = None, + prev_module_hook: Optional[UserCpuOffloadHook] = None, +): + """ + Offloads a model on the CPU and puts it back to an execution device when executed. The difference with + [`cpu_offload`] is that the model stays on the execution device after the forward and is only offloaded again when + the `offload` method of the returned `hook` is called. Useful for pipelines running a model in a loop. + + Args: + model (`torch.nn.Module`): + The model to offload. + execution_device(`str`, `int` or `torch.device`, *optional*): + The device on which the model should be executed. Will default to the MPS device if it's available, then + GPU 0 if there is a GPU, and finally to the CPU. + prev_module_hook (`UserCpuOffloadHook`, *optional*): + The hook sent back by this function for a previous model in the pipeline you are running. If passed, its + offload method will be called just before the forward of the model to which this hook is attached. + + Example: + + ```py + model_1, hook_1 = cpu_offload_with_hook(model_1, cuda_device) + model_2, hook_2 = cpu_offload_with_hook(model_2, cuda_device, prev_module_hook=hook_1) + model_3, hook_3 = cpu_offload_with_hook(model_3, cuda_device, prev_module_hook=hook_2) + + hid_1 = model_1(input) + for i in range(50): + # model1 is offloaded on the CPU at the first iteration, model 2 stays on the GPU for this whole loop. + hid_2 = model_2(hid_1) + # model2 is offloaded to the CPU just before this forward. + hid_3 = model_3(hid_3) + + # For model3, you need to manually call the hook offload method. + hook_3.offload() + ``` + """ + hook = CpuOffload(execution_device=execution_device, prev_module_hook=prev_module_hook) + add_hook_to_module(model, hook, append=True) + user_hook = UserCpuOffloadHook(model, hook) + return model, user_hook + + +def disk_offload( + model: nn.Module, + offload_dir: Union[str, os.PathLike], + execution_device: Optional[torch.device] = None, + offload_buffers: bool = False, + preload_module_classes: Optional[list[str]] = None, +): + """ + Activates full disk offload for a model. As a result, all parameters of the model will be offloaded as + memory-mapped array in a given folder. During the forward pass, parameters will be accessed from that folder and + put on the execution device passed as they are needed, then offloaded again. + + Args: + model (`torch.nn.Module`): The model to offload. + offload_dir (`str` or `os.PathLike`): + The folder in which to offload the model weights (or where the model weights are already offloaded). + execution_device (`torch.device`, *optional*): + The device on which the forward pass of the model will be executed (should be a GPU). Will default to the + model's first parameter device. + offload_buffers (`bool`, *optional*, defaults to `False`): + Whether or not to offload the buffers with the model parameters. + preload_module_classes (`List[str]`, *optional*): + A list of classes whose instances should load all their weights (even in the submodules) at the beginning + of the forward. This should only be used for classes that have submodules which are registered but not + called directly during the forward, for instance if a `dense` linear layer is registered, but at forward, + `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly. + """ + if not os.path.isdir(offload_dir) or not os.path.isfile(os.path.join(offload_dir, "index.json")): + offload_state_dict(offload_dir, model.state_dict()) + if execution_device is None: + execution_device = next(iter(model.parameters())).device + weights_map = OffloadedWeightsLoader(save_folder=offload_dir) + + add_hook_to_module(model, AlignDevicesHook(io_same_device=True), append=True) + attach_align_device_hook( + model, + execution_device=execution_device, + offload=True, + offload_buffers=offload_buffers, + weights_map=weights_map, + preload_module_classes=preload_module_classes, + ) + + return model + + +def dispatch_model( + model: nn.Module, + device_map: dict[str, Union[str, int, torch.device]], + main_device: Optional[torch.device] = None, + state_dict: Optional[dict[str, torch.Tensor]] = None, + offload_dir: Optional[Union[str, os.PathLike]] = None, + offload_index: Optional[dict[str, str]] = None, + offload_buffers: bool = False, + skip_keys: Optional[Union[str, list[str]]] = None, + preload_module_classes: Optional[list[str]] = None, + force_hooks: bool = False, +): + """ + Dispatches a model according to a given device map. Layers of the model might be spread across GPUs, offloaded on + the CPU or even the disk. + + Args: + model (`torch.nn.Module`): + The model to dispatch. + device_map (`Dict[str, Union[str, int, torch.device]]`): + A dictionary mapping module names in the models `state_dict` to the device they should go to. Note that + `"disk"` is accepted even if it's not a proper value for `torch.device`. + main_device (`str`, `int` or `torch.device`, *optional*): + The main execution device. Will default to the first device in the `device_map` different from `"cpu"` or + `"disk"`. + state_dict (`Dict[str, torch.Tensor]`, *optional*): + The state dict of the part of the model that will be kept on CPU. + offload_dir (`str` or `os.PathLike`): + The folder in which to offload the model weights (or where the model weights are already offloaded). + offload_index (`Dict`, *optional*): + A dictionary from weight name to their information (`dtype`/ `shape` or safetensors filename). Will default + to the index saved in `save_folder`. + offload_buffers (`bool`, *optional*, defaults to `False`): + Whether or not to offload the buffers with the model parameters. + skip_keys (`str` or `List[str]`, *optional*): + A list of keys to ignore when moving inputs or outputs between devices. + preload_module_classes (`List[str]`, *optional*): + A list of classes whose instances should load all their weights (even in the submodules) at the beginning + of the forward. This should only be used for classes that have submodules which are registered but not + called directly during the forward, for instance if a `dense` linear layer is registered, but at forward, + `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly. + force_hooks (`bool`, *optional*, defaults to `False`): + Whether or not to force device hooks to be attached to the model even if all layers are dispatched to a + single device. + """ + # Error early if the device map is incomplete. + check_device_map(model, device_map) + + # We need to force hook for quantized model that can't be moved with to() + if getattr(model, "quantization_method", "bitsandbytes") == "bitsandbytes": + # since bnb 0.43.2, we can move 4-bit model + if getattr(model, "is_loaded_in_8bit", False) or ( + getattr(model, "is_loaded_in_4bit", False) and not is_bnb_available(min_version="0.43.2") + ): + force_hooks = True + + # We attach hooks if the device_map has at least 2 different devices or if + # force_hooks is set to `True`. Otherwise, the model in already loaded + # in the unique device and the user can decide where to dispatch the model. + # If the model is quantized, we always force-dispatch the model + if (len(set(device_map.values())) > 1) or force_hooks: + if main_device is None: + if set(device_map.values()) == {"cpu"} or set(device_map.values()) == {"cpu", "disk"}: + main_device = "cpu" + else: + main_device = [d for d in device_map.values() if d not in ["cpu", "disk"]][0] + + if main_device != "cpu": + cpu_modules = [name for name, device in device_map.items() if device == "cpu"] + if state_dict is None and len(cpu_modules) > 0: + state_dict = extract_submodules_state_dict(model.state_dict(), cpu_modules) + + disk_modules = [name for name, device in device_map.items() if device == "disk"] + if offload_dir is None and offload_index is None and len(disk_modules) > 0: + raise ValueError( + "We need an `offload_dir` to dispatch this model according to this `device_map`, the following submodules " + f"need to be offloaded: {', '.join(disk_modules)}." + ) + if ( + len(disk_modules) > 0 + and offload_index is None + and (not os.path.isdir(offload_dir) or not os.path.isfile(os.path.join(offload_dir, "index.json"))) + ): + disk_state_dict = extract_submodules_state_dict(model.state_dict(), disk_modules) + offload_state_dict(offload_dir, disk_state_dict) + + execution_device = { + name: main_device if device in ["cpu", "disk"] else device for name, device in device_map.items() + } + execution_device[""] = main_device + offloaded_devices = ["disk"] if main_device == "cpu" or main_device == "mps" else ["cpu", "disk"] + offload = {name: device in offloaded_devices for name, device in device_map.items()} + save_folder = offload_dir if len(disk_modules) > 0 else None + if state_dict is not None or save_folder is not None or offload_index is not None: + device = main_device if offload_index is not None else None + weights_map = OffloadedWeightsLoader( + state_dict=state_dict, save_folder=save_folder, index=offload_index, device=device + ) + else: + weights_map = None + + # When dispatching the model's parameters to the devices specified in device_map, we want to avoid allocating memory several times for the + # tied parameters. The dictionary tied_params_map keeps track of the already allocated data for a given tied parameter (represented by its + # original pointer) on each devices. + tied_params = find_tied_parameters(model) + + tied_params_map = {} + for group in tied_params: + for param_name in group: + # data_ptr() is enough here, as `find_tied_parameters` finds tied params simply by comparing `param1 is param2`, so we don't need + # to care about views of tensors through storage_offset. + data_ptr = recursive_getattr(model, param_name).data_ptr() + tied_params_map[data_ptr] = {} + + # Note: To handle the disk offloading case, we can not simply use weights_map[param_name].data_ptr() as the reference pointer, + # as we have no guarantee that safetensors' `file.get_tensor()` will always give the same pointer. + + attach_align_device_hook_on_blocks( + model, + execution_device=execution_device, + offload=offload, + offload_buffers=offload_buffers, + weights_map=weights_map, + skip_keys=skip_keys, + preload_module_classes=preload_module_classes, + tied_params_map=tied_params_map, + ) + + # warn if there is any params on the meta device + offloaded_devices_str = " and ".join( + [device for device in set(device_map.values()) if device in ("cpu", "disk")] + ) + if len(offloaded_devices_str) > 0: + logger.warning( + f"Some parameters are on the meta device because they were offloaded to the {offloaded_devices_str}." + ) + + # Attaching the hook may break tied weights, so we retie them + retie_parameters(model, tied_params) + + # add warning to cuda and to method + def add_warning(fn, model): + @wraps(fn) + def wrapper(*args, **kwargs): + warning_msg = "You shouldn't move a model that is dispatched using accelerate hooks." + if str(fn.__name__) == "to": + to_device = torch._C._nn._parse_to(*args, **kwargs)[0] + if to_device is not None: + logger.warning(warning_msg) + else: + logger.warning(warning_msg) + for param in model.parameters(): + if param.device == torch.device("meta"): + raise RuntimeError("You can't move a model that has some modules offloaded to cpu or disk.") + return fn(*args, **kwargs) + + return wrapper + + # Make sure to update _accelerate_added_attributes in hooks.py if you add any hook + model.to = add_warning(model.to, model) + if is_npu_available(): + model.npu = add_warning(model.npu, model) + elif is_mlu_available(): + model.mlu = add_warning(model.mlu, model) + elif is_sdaa_available(): + model.sdaa = add_warning(model.sdaa, model) + elif is_musa_available(): + model.musa = add_warning(model.musa, model) + elif is_xpu_available(): + model.xpu = add_warning(model.xpu, model) + else: + model.cuda = add_warning(model.cuda, model) + + # Check if we are using multi-gpus with RTX 4000 series + use_multi_gpu = len([device for device in set(device_map.values()) if device not in ("cpu", "disk")]) > 1 + if use_multi_gpu and not check_cuda_p2p_ib_support(): + logger.warning( + "We've detected an older driver with an RTX 4000 series GPU. These drivers have issues with P2P. " + "This can affect the multi-gpu inference when using accelerate device_map." + "Please make sure to update your driver to the latest version which resolves this." + ) + else: + device = list(device_map.values())[0] + # `torch.Tensor.to()` is not supported by `torch_npu` (see this [issue](https://github.com/Ascend/pytorch/issues/16)). + if is_npu_available() and isinstance(device, int): + device = f"npu:{device}" + elif is_mlu_available() and isinstance(device, int): + device = f"mlu:{device}" + elif is_sdaa_available() and isinstance(device, int): + device = f"sdaa:{device}" + elif is_musa_available() and isinstance(device, int): + device = f"musa:{device}" + if device != "disk": + model.to(device) + else: + raise ValueError( + "You are trying to offload the whole model to the disk. Please use the `disk_offload` function instead." + ) + # Convert OrderedDict back to dict for easier usage + model.hf_device_map = dict(device_map) + return model + + +def load_checkpoint_and_dispatch( + model: nn.Module, + checkpoint: Union[str, os.PathLike], + device_map: Optional[Union[str, dict[str, Union[int, str, torch.device]]]] = None, + max_memory: Optional[dict[Union[int, str], Union[int, str]]] = None, + no_split_module_classes: Optional[list[str]] = None, + offload_folder: Optional[Union[str, os.PathLike]] = None, + offload_buffers: bool = False, + dtype: Optional[Union[str, torch.dtype]] = None, + offload_state_dict: Optional[bool] = None, + skip_keys: Optional[Union[str, list[str]]] = None, + preload_module_classes: Optional[list[str]] = None, + force_hooks: bool = False, + strict: bool = False, + full_state_dict: bool = True, + broadcast_from_rank0: bool = False, +): + """ + Loads a (potentially sharded) checkpoint inside a model, potentially sending weights to a given device as they are + loaded and adds the various hooks that will make this model run properly (even if split across devices). + + Args: + model (`torch.nn.Module`): The model in which we want to load a checkpoint. + checkpoint (`str` or `os.PathLike`): + The folder checkpoint to load. It can be: + - a path to a file containing a whole model state dict + - a path to a `.json` file containing the index to a sharded checkpoint + - a path to a folder containing a unique `.index.json` file and the shards of a checkpoint. + device_map (`Dict[str, Union[int, str, torch.device]]`, *optional*): + A map that specifies where each submodule should go. It doesn't need to be refined to each parameter/buffer + name, once a given module name is inside, every submodule of it will be sent to the same device. + + To have Accelerate compute the most optimized `device_map` automatically, set `device_map="auto"`. For more + information about each option see [here](../concept_guides/big_model_inference#designing-a-device-map). + Defaults to None, which means [`dispatch_model`] will not be called. + max_memory (`Dict`, *optional*): + A dictionary device identifier to maximum memory. Will default to the maximum memory available for each GPU + and the available CPU RAM if unset. + no_split_module_classes (`List[str]`, *optional*): + A list of layer class names that should never be split across device (for instance any layer that has a + residual connection). + offload_folder (`str` or `os.PathLike`, *optional*): + If the `device_map` contains any value `"disk"`, the folder where we will offload weights. + offload_buffers (`bool`, *optional*, defaults to `False`): + In the layers that are offloaded on the CPU or the hard drive, whether or not to offload the buffers as + well as the parameters. + dtype (`str` or `torch.dtype`, *optional*): + If provided, the weights will be converted to that type when loaded. + offload_state_dict (`bool`, *optional*): + If `True`, will temporarily offload the CPU state dict on the hard drive to avoid getting out of CPU RAM if + the weight of the CPU state dict + the biggest shard does not fit. Will default to `True` if the device map + picked contains `"disk"` values. + skip_keys (`str` or `List[str]`, *optional*): + A list of keys to ignore when moving inputs or outputs between devices. + preload_module_classes (`List[str]`, *optional*): + A list of classes whose instances should load all their weights (even in the submodules) at the beginning + of the forward. This should only be used for classes that have submodules which are registered but not + called directly during the forward, for instance if a `dense` linear layer is registered, but at forward, + `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly. + force_hooks (`bool`, *optional*, defaults to `False`): + Whether or not to force device hooks to be attached to the model even if all layers are dispatched to a + single device. + strict (`bool`, *optional*, defaults to `False`): + Whether to strictly enforce that the keys in the checkpoint state_dict match the keys of the model's + state_dict. + full_state_dict (`bool`, *optional*, defaults to `True`): if this is set to `True`, all the tensors in the + loaded state_dict will be gathered. No ShardedTensor and DTensor will be in the loaded state_dict. + broadcast_from_rank0 (`False`, *optional*, defaults to `False`): when the option is `True`, a distributed + `ProcessGroup` must be initialized. rank0 should receive a full state_dict and will broadcast the tensors + in the state_dict one by one to other ranks. Other ranks will receive the tensors and shard (if applicable) + according to the local shards in the model. + + Example: + + ```python + >>> from accelerate import init_empty_weights, load_checkpoint_and_dispatch + >>> from huggingface_hub import hf_hub_download + >>> from transformers import AutoConfig, AutoModelForCausalLM + + >>> # Download the Weights + >>> checkpoint = "EleutherAI/gpt-j-6B" + >>> weights_location = hf_hub_download(checkpoint, "pytorch_model.bin") + + >>> # Create a model and initialize it with empty weights + >>> config = AutoConfig.from_pretrained(checkpoint) + >>> with init_empty_weights(): + ... model = AutoModelForCausalLM.from_config(config) + + >>> # Load the checkpoint and dispatch it to the right devices + >>> model = load_checkpoint_and_dispatch( + ... model, weights_location, device_map="auto", no_split_module_classes=["GPTJBlock"] + ... ) + ``` + """ + if isinstance(device_map, str) and device_map not in ["auto", "balanced", "balanced_low_0", "sequential"]: + raise ValueError( + "If passing a string for `device_map`, please choose 'auto', 'balanced', 'balanced_low_0' or 'sequential'." + ) + if isinstance(device_map, str): + if device_map != "sequential": + max_memory = get_balanced_memory( + model, + max_memory=max_memory, + no_split_module_classes=no_split_module_classes, + dtype=dtype, + low_zero=(device_map == "balanced_low_0"), + ) + device_map = infer_auto_device_map( + model, + max_memory=max_memory, + no_split_module_classes=no_split_module_classes, + dtype=dtype, + offload_buffers=offload_buffers, + ) + if offload_state_dict is None and device_map is not None and "disk" in device_map.values(): + offload_state_dict = True + load_checkpoint_in_model( + model, + checkpoint, + device_map=device_map, + offload_folder=offload_folder, + dtype=dtype, + offload_state_dict=offload_state_dict, + offload_buffers=offload_buffers, + strict=strict, + full_state_dict=full_state_dict, + broadcast_from_rank0=broadcast_from_rank0, + ) + if device_map is None: + return model + return dispatch_model( + model, + device_map=device_map, + offload_dir=offload_folder, + offload_buffers=offload_buffers, + skip_keys=skip_keys, + preload_module_classes=preload_module_classes, + force_hooks=force_hooks, + ) + + +def attach_layerwise_casting_hooks( + module: torch.nn.Module, + storage_dtype: torch.dtype, + compute_dtype: torch.dtype, + skip_modules_pattern: Optional[Union[str, tuple[str, ...]]] = None, + skip_modules_classes: Optional[tuple[type[torch.nn.Module], ...]] = None, + non_blocking: bool = False, +) -> None: + r""" + Applies layerwise casting to a given module. The module expected here is a PyTorch `nn.Module`. This is helpful for + reducing memory requirements when one doesn't want to fully quantize a model. Model params can be kept in say, + `torch.float8_e4m3fn` and upcasted to a higher precision like `torch.bfloat16` during forward pass and downcasted + back to `torch.float8_e4m3fn` to realize memory savings. + + Args: + module (`torch.nn.Module`): + The module whose leaf modules will be cast to a high precision dtype for computation, and to a low + precision dtype for storage. + storage_dtype (`torch.dtype`): + The dtype to cast the module to before/after the forward pass for storage. + compute_dtype (`torch.dtype`): + The dtype to cast the module to during the forward pass for computation. + skip_modules_pattern (`tuple[str, ...]`, defaults to `None`): + A list of patterns to match the names of the modules to skip during the layerwise casting process. If set + to `None` alongside `skip_modules_classes` being `None`, the layerwise casting is applied directly to the + module instead of its internal submodules. + skip_modules_classes (`tuple[type[torch.nn.Module], ...]`, defaults to `None`): + A list of module classes to skip during the layerwise casting process. + non_blocking (`bool`, defaults to `False`): + If `True`, the weight casting operations are non-blocking. + + Example: + + ```python + >>> from accelerate.hooks import attach_layerwise_casting_hooks + >>> from transformers import AutoModelForCausalLM + >>> import torch + + >>> # Model + >>> checkpoint = "EleutherAI/gpt-j-6B" + >>> model = AutoModelForCausalLM.from_pretrained(checkpoint) + + >>> # Attach hooks and perform inference + >>> attach_layerwise_casting_hooks(model, storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16) + >>> with torch.no_grad(): + ... model(...) + ``` + + Users can also pass modules they want to avoid from getting downcasted. + + ```py + >>> attach_layerwise_casting_hooks( + ... model, storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16, skip_modules_pattern=["norm"] + ... ) + ``` + """ + _attach_layerwise_casting_hooks( + module, storage_dtype, compute_dtype, skip_modules_pattern, skip_modules_classes, non_blocking + ) + + +def _attach_layerwise_casting_hooks( + module: torch.nn.Module, + storage_dtype: torch.dtype, + compute_dtype: torch.dtype, + skip_modules_pattern: Optional[Union[str, tuple[str, ...]]] = None, + skip_modules_classes: Optional[tuple[type[torch.nn.Module], ...]] = None, + non_blocking: bool = False, + _prefix: str = "", +): + should_skip = (skip_modules_classes is not None and isinstance(module, skip_modules_classes)) or ( + skip_modules_pattern is not None and any(re.search(pattern, _prefix) for pattern in skip_modules_pattern) + ) + if should_skip: + logger.debug(f'Skipping layerwise casting for layer "{_prefix}"') + return + + if isinstance(module, SUPPORTED_PYTORCH_LAYERS_FOR_UPCASTING): + logger.debug(f'Applying layerwise casting to layer "{_prefix}"') + add_hook_to_module( + module, + LayerwiseCastingHook(storage_dtype=storage_dtype, compute_dtype=compute_dtype, non_blocking=non_blocking), + append=True, + ) + return + + for name, submodule in module.named_children(): + layer_name = f"{_prefix}.{name}" if _prefix else name + _attach_layerwise_casting_hooks( + submodule, + storage_dtype, + compute_dtype, + skip_modules_pattern, + skip_modules_classes, + non_blocking, + _prefix=layer_name, + ) + + +def _attach_context_parallel_hooks( + model: nn.Module, +): + """ + Monkeypatch huggingface's `transformers` model to fix attention mask issues when using context parallelism. + + This function attaches forward_pre_hooks to each self_attn module of the model, where each hook checks the + args/kwargs, if they contain an attention mask, if it does, it will remove this mask, check if it is a causal mask, + if yes, will add a kwarg `is_causal=True`, otherwise will raise an error. This is because context parallelism does + not support attention masks. This function modifies the model in place. + + Args: + model (`nn.Module`): + The model to attach the hooks to. + + """ + + def _self_attn_pre_forward_hook(_module, module_args, module_kwargs): + if "attention_mask" in module_kwargs: + module_kwargs["attention_mask"] = None + module_kwargs["is_causal"] = True + + return module_args, module_kwargs + + for name, module in model.named_modules(): + # We hope (assume) that if user uses their own model (without this structure which transformers uses), they read the docs saying they can't pass in attention masks + # Then these cases can happen: + # 1) some modules end with a `self-attn` module, in which case we attach the hook, but the + # there's no attention mask kwarg -> hook is a no-op + # 2) some modules end with a `self-attn` module, in which case we attach the hook, and the + # attention mask kwarg is passed -> hook will remove the attention mask and add + # `is_causal=True` kwarg, which either crashes the training or fixes it + # (training would crash anyway as attention mask isn't supported) + # 3) no modules end with a `self-attn` module, in which case we don't attach the hook, this is + # a no-op as well + if name.endswith("self_attn"): + # we want the hook to be executed first, to avoid any other hooks doing work on the attention mask + module.register_forward_pre_hook(_self_attn_pre_forward_hook, with_kwargs=True, prepend=True) diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/checkpointing.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/checkpointing.py new file mode 100644 index 0000000000000000000000000000000000000000..57e58300a6f579451f79cf8bd0c26cbbb65de04c --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/checkpointing.py @@ -0,0 +1,331 @@ +# Copyright 2022 The HuggingFace Team. All rights reserved. +# +# 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. + +import random +from pathlib import Path +from typing import Optional + +import numpy as np +import torch +from safetensors.torch import load_model + +from .utils import ( + MODEL_NAME, + OPTIMIZER_NAME, + RNG_STATE_NAME, + SAFE_MODEL_NAME, + SAFE_WEIGHTS_NAME, + SAMPLER_NAME, + SCALER_NAME, + SCHEDULER_NAME, + WEIGHTS_NAME, + get_pretty_name, + is_cuda_available, + is_hpu_available, + is_mlu_available, + is_musa_available, + is_sdaa_available, + is_torch_version, + is_torch_xla_available, + is_xpu_available, + load, + save, +) + + +if is_torch_version(">=", "2.4.0"): + from torch.amp import GradScaler +else: + from torch.cuda.amp import GradScaler + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + +from .logging import get_logger +from .state import PartialState + + +logger = get_logger(__name__) + + +def save_accelerator_state( + output_dir: str, + model_states: list[dict], + optimizers: list, + schedulers: list, + dataloaders: list, + process_index: int, + step: int, + scaler: Optional[GradScaler] = None, + save_on_each_node: bool = False, + safe_serialization: bool = True, +): + """ + Saves the current states of the models, optimizers, scaler, and RNG generators to a given directory. + + + + If `safe_serialization` is `True`, models will be saved with `safetensors` while the rest are saved using native + `pickle`. + + + + Args: + output_dir (`str` or `os.PathLike`): + The name of the folder to save all relevant weights and states. + model_states (`List[torch.nn.Module]`): + A list of model states + optimizers (`List[torch.optim.Optimizer]`): + A list of optimizer instances + schedulers (`List[torch.optim.lr_scheduler._LRScheduler]`): + A list of learning rate schedulers + dataloaders (`List[torch.utils.data.DataLoader]`): + A list of dataloader instances to save their sampler states + process_index (`int`): + The current process index in the Accelerator state + step (`int`): + The current step in the internal step tracker + scaler (`torch.amp.GradScaler`, *optional*): + An optional gradient scaler instance to save; + save_on_each_node (`bool`, *optional*): + Whether to save on every node, or only the main node. + safe_serialization (`bool`, *optional*, defaults to `True`): + Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`). + """ + output_dir = Path(output_dir) + # Model states + for i, state in enumerate(model_states): + weights_name = WEIGHTS_NAME if not safe_serialization else SAFE_WEIGHTS_NAME + if i > 0: + weights_name = weights_name.replace(".", f"_{i}.") + output_model_file = output_dir.joinpath(weights_name) + save(state, output_model_file, save_on_each_node=save_on_each_node, safe_serialization=safe_serialization) + logger.info(f"Model weights saved in {output_model_file}") + # Optimizer states + for i, opt in enumerate(optimizers): + state = opt.state_dict() + optimizer_name = f"{OPTIMIZER_NAME}.bin" if i == 0 else f"{OPTIMIZER_NAME}_{i}.bin" + output_optimizer_file = output_dir.joinpath(optimizer_name) + save(state, output_optimizer_file, save_on_each_node=save_on_each_node, safe_serialization=False) + logger.info(f"Optimizer state saved in {output_optimizer_file}") + # Scheduler states + for i, scheduler in enumerate(schedulers): + state = scheduler.state_dict() + scheduler_name = f"{SCHEDULER_NAME}.bin" if i == 0 else f"{SCHEDULER_NAME}_{i}.bin" + output_scheduler_file = output_dir.joinpath(scheduler_name) + save(state, output_scheduler_file, save_on_each_node=save_on_each_node, safe_serialization=False) + logger.info(f"Scheduler state saved in {output_scheduler_file}") + # DataLoader states + for i, dataloader in enumerate(dataloaders): + sampler_name = f"{SAMPLER_NAME}.bin" if i == 0 else f"{SAMPLER_NAME}_{i}.bin" + output_sampler_file = output_dir.joinpath(sampler_name) + # Only save if we have our custom sampler + from .data_loader import IterableDatasetShard, SeedableRandomSampler + + if isinstance(dataloader.dataset, IterableDatasetShard): + sampler = dataloader.get_sampler() + if isinstance(sampler, SeedableRandomSampler): + save(sampler, output_sampler_file, save_on_each_node=save_on_each_node, safe_serialization=False) + if getattr(dataloader, "use_stateful_dataloader", False): + dataloader_state_dict_name = "dl_state_dict.bin" if i == 0 else f"dl_state_dict_{i}.bin" + output_dataloader_state_dict_file = output_dir.joinpath(dataloader_state_dict_name) + state_dict = dataloader.state_dict() + torch.save(state_dict, output_dataloader_state_dict_file) + logger.info(f"Sampler state for dataloader {i} saved in {output_sampler_file}") + + # GradScaler state + if scaler is not None: + state = scaler.state_dict() + output_scaler_file = output_dir.joinpath(SCALER_NAME) + torch.save(state, output_scaler_file) + logger.info(f"Gradient scaler state saved in {output_scaler_file}") + # Random number generator states + states = {} + states_name = f"{RNG_STATE_NAME}_{process_index}.pkl" + states["step"] = step + states["random_state"] = random.getstate() + states["numpy_random_seed"] = np.random.get_state() + states["torch_manual_seed"] = torch.get_rng_state() + if is_xpu_available(): + states["torch_xpu_manual_seed"] = torch.xpu.get_rng_state_all() + if is_mlu_available(): + states["torch_mlu_manual_seed"] = torch.mlu.get_rng_state_all() + elif is_sdaa_available(): + states["torch_sdaa_manual_seed"] = torch.sdaa.get_rng_state_all() + elif is_musa_available(): + states["torch_musa_manual_seed"] = torch.musa.get_rng_state_all() + if is_hpu_available(): + states["torch_hpu_manual_seed"] = torch.hpu.get_rng_state_all() + if is_cuda_available(): + states["torch_cuda_manual_seed"] = torch.cuda.get_rng_state_all() + if is_torch_xla_available(): + states["xm_seed"] = xm.get_rng_state() + output_states_file = output_dir.joinpath(states_name) + torch.save(states, output_states_file) + logger.info(f"Random states saved in {output_states_file}") + return output_dir + + +def load_accelerator_state( + input_dir, + models, + optimizers, + schedulers, + dataloaders, + process_index, + scaler=None, + map_location=None, + load_kwargs=None, + **load_model_func_kwargs, +): + """ + Loads states of the models, optimizers, scaler, and RNG generators from a given directory. + + Args: + input_dir (`str` or `os.PathLike`): + The name of the folder to load all relevant weights and states. + models (`List[torch.nn.Module]`): + A list of model instances + optimizers (`List[torch.optim.Optimizer]`): + A list of optimizer instances + schedulers (`List[torch.optim.lr_scheduler._LRScheduler]`): + A list of learning rate schedulers + process_index (`int`): + The current process index in the Accelerator state + scaler (`torch.amp.GradScaler`, *optional*): + An optional *GradScaler* instance to load + map_location (`str`, *optional*): + What device to load the optimizer state onto. Should be one of either "cpu" or "on_device". + load_kwargs (`dict`, *optional*): + Additional arguments that can be passed to the `load` function. + load_model_func_kwargs (`dict`, *optional*): + Additional arguments that can be passed to the model's `load_state_dict` method. + + Returns: + `dict`: Contains the `Accelerator` attributes to override while loading the state. + """ + # stores the `Accelerator` attributes to override + override_attributes = dict() + if map_location not in [None, "cpu", "on_device"]: + raise TypeError( + "Unsupported optimizer map location passed, please choose one of `None`, `'cpu'`, or `'on_device'`" + ) + if map_location is None: + map_location = "cpu" + elif map_location == "on_device": + map_location = PartialState().device + + if load_kwargs is None: + load_kwargs = {} + + input_dir = Path(input_dir) + # Model states + for i, model in enumerate(models): + ending = f"_{i}" if i > 0 else "" + input_model_file = input_dir.joinpath(f"{SAFE_MODEL_NAME}{ending}.safetensors") + if input_model_file.exists(): + load_model(model, input_model_file, device=str(map_location), **load_model_func_kwargs) + else: + # Load with torch + input_model_file = input_dir.joinpath(f"{MODEL_NAME}{ending}.bin") + state_dict = load(input_model_file, map_location=map_location) + model.load_state_dict(state_dict, **load_model_func_kwargs) + logger.info("All model weights loaded successfully") + + # Optimizer states + for i, opt in enumerate(optimizers): + optimizer_name = f"{OPTIMIZER_NAME}.bin" if i == 0 else f"{OPTIMIZER_NAME}_{i}.bin" + input_optimizer_file = input_dir.joinpath(optimizer_name) + optimizer_state = load(input_optimizer_file, map_location=map_location, **load_kwargs) + optimizers[i].load_state_dict(optimizer_state) + logger.info("All optimizer states loaded successfully") + + # Scheduler states + for i, scheduler in enumerate(schedulers): + scheduler_name = f"{SCHEDULER_NAME}.bin" if i == 0 else f"{SCHEDULER_NAME}_{i}.bin" + input_scheduler_file = input_dir.joinpath(scheduler_name) + scheduler_state = load(input_scheduler_file, **load_kwargs) + scheduler.load_state_dict(scheduler_state) + logger.info("All scheduler states loaded successfully") + + for i, dataloader in enumerate(dataloaders): + sampler_name = f"{SAMPLER_NAME}.bin" if i == 0 else f"{SAMPLER_NAME}_{i}.bin" + input_sampler_file = input_dir.joinpath(sampler_name) + # Only load if we have our custom sampler + from .data_loader import IterableDatasetShard, SeedableRandomSampler + + if isinstance(dataloader.dataset, IterableDatasetShard): + sampler = dataloader.get_sampler() + if isinstance(sampler, SeedableRandomSampler): + sampler = dataloader.set_sampler(load(input_sampler_file)) + if getattr(dataloader, "use_stateful_dataloader", False): + dataloader_state_dict_name = "dl_state_dict.bin" if i == 0 else f"dl_state_dict_{i}.bin" + input_dataloader_state_dict_file = input_dir.joinpath(dataloader_state_dict_name) + if input_dataloader_state_dict_file.exists(): + state_dict = load(input_dataloader_state_dict_file, **load_kwargs) + dataloader.load_state_dict(state_dict) + logger.info("All dataloader sampler states loaded successfully") + + # GradScaler state + if scaler is not None: + input_scaler_file = input_dir.joinpath(SCALER_NAME) + scaler_state = load(input_scaler_file) + scaler.load_state_dict(scaler_state) + logger.info("GradScaler state loaded successfully") + + # Random states + try: + states = load(input_dir.joinpath(f"{RNG_STATE_NAME}_{process_index}.pkl")) + if "step" in states: + override_attributes["step"] = states["step"] + random.setstate(states["random_state"]) + np.random.set_state(states["numpy_random_seed"]) + torch.set_rng_state(states["torch_manual_seed"]) + if is_xpu_available(): + torch.xpu.set_rng_state_all(states["torch_xpu_manual_seed"]) + if is_mlu_available(): + torch.mlu.set_rng_state_all(states["torch_mlu_manual_seed"]) + elif is_sdaa_available(): + torch.sdaa.set_rng_state_all(states["torch_sdaa_manual_seed"]) + elif is_musa_available(): + torch.musa.set_rng_state_all(states["torch_musa_manual_seed"]) + else: + torch.cuda.set_rng_state_all(states["torch_cuda_manual_seed"]) + if is_torch_xla_available(): + xm.set_rng_state(states["xm_seed"]) + logger.info("All random states loaded successfully") + except Exception: + logger.info("Could not load random states") + + return override_attributes + + +def save_custom_state(obj, path, index: int = 0, save_on_each_node: bool = False): + """ + Saves the state of `obj` to `{path}/custom_checkpoint_{index}.pkl` + """ + # Should this be the right way to get a qual_name type value from `obj`? + save_location = Path(path) / f"custom_checkpoint_{index}.pkl" + logger.info(f"Saving the state of {get_pretty_name(obj)} to {save_location}") + save(obj.state_dict(), save_location, save_on_each_node=save_on_each_node) + + +def load_custom_state(obj, path, index: int = 0): + """ + Loads the state of `obj` at `{path}/custom_checkpoint_{index}.pkl`. Will always set `weights_only=False` when + loading the state. + """ + load_location = f"{path}/custom_checkpoint_{index}.pkl" + logger.info(f"Loading the state of {get_pretty_name(obj)} from {load_location}") + obj.load_state_dict(load(load_location, map_location="cpu", weights_only=False)) diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/data_loader.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/data_loader.py new file mode 100644 index 0000000000000000000000000000000000000000..e06f9fddf146cc9ddc935a08d4e70d9cf02a830a --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/data_loader.py @@ -0,0 +1,1451 @@ +# Copyright 2021 The HuggingFace Team. All rights reserved. +# +# 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. + +import importlib +import math +from contextlib import suppress +from typing import Callable, Optional, Union + +import torch +from packaging import version +from torch.utils.data import BatchSampler, DataLoader, IterableDataset, RandomSampler + +from .logging import get_logger +from .state import DistributedType, GradientState, PartialState, is_torch_xla_available +from .utils import ( + RNGType, + broadcast, + broadcast_object_list, + compare_versions, + concatenate, + find_batch_size, + get_data_structure, + initialize_tensors, + is_datasets_available, + is_torch_version, + is_torchdata_stateful_dataloader_available, + send_to_device, + slice_tensors, + synchronize_rng_states, +) + + +logger = get_logger(__name__) + +# kwargs of the DataLoader in min version 2.0 +_PYTORCH_DATALOADER_KWARGS = { + "batch_size": 1, + "shuffle": False, + "sampler": None, + "batch_sampler": None, + "num_workers": 0, + "collate_fn": None, + "pin_memory": False, + "drop_last": False, + "timeout": 0, + "worker_init_fn": None, + "multiprocessing_context": None, + "generator": None, + "prefetch_factor": 2, + "persistent_workers": False, + "pin_memory_device": "", +} + +# kwargs added after by version +_PYTORCH_DATALOADER_ADDITIONAL_KWARGS = {"2.6.0": {"in_order": True}} + +for v, additional_kwargs in _PYTORCH_DATALOADER_ADDITIONAL_KWARGS.items(): + if is_torch_version(">=", v): + _PYTORCH_DATALOADER_KWARGS.update(additional_kwargs) + + +class SeedableRandomSampler(RandomSampler): + """ + Same as a random sampler, except that in `__iter__` a seed can be used. + + Needed specifically in distributed cases, when the random generator for each GPU needs to start from the same seed + and be fully reproducible on multiple iterations. + + If a custom `generator` is passed, it will rely on its initial seed as well as the current iteration it is on + (stored in `self.epoch`). + """ + + def __init__(self, *args, **kwargs): + data_seed = kwargs.pop("data_seed", None) + super().__init__(*args, **kwargs) + + self.initial_seed = data_seed if data_seed is not None else torch.random.initial_seed() + self.epoch = 0 + + def __iter__(self): + if self.generator is None: + self.generator = torch.Generator( + device=torch.get_default_device() if hasattr(torch, "get_default_device") else "cpu" + ) + self.generator.manual_seed(self.initial_seed) + + # Allow `self.epoch` to modify the seed of the generator + seed = self.epoch + self.initial_seed + # print("Setting seed at epoch", self.epoch, seed) + self.generator.manual_seed(seed) + yield from super().__iter__() + self.set_epoch(self.epoch + 1) + + def set_epoch(self, epoch: int): + "Sets the current iteration of the sampler." + self.epoch = epoch + + +class BatchSamplerShard(BatchSampler): + """ + Wraps a PyTorch `BatchSampler` to generate batches for one of the processes only. Instances of this class will + always yield a number of batches that is a round multiple of `num_processes` and that all have the same size. + Depending on the value of the `drop_last` attribute of the batch sampler passed, it will either stop the iteration + at the first batch that would be too small / not present on all processes or loop with indices from the beginning. + + Args: + batch_sampler (`torch.utils.data.sampler.BatchSampler`): + The batch sampler to split in several shards. + num_processes (`int`, *optional*, defaults to 1): + The number of processes running concurrently. + process_index (`int`, *optional*, defaults to 0): + The index of the current process. + split_batches (`bool`, *optional*, defaults to `False`): + Whether the shards should be created by splitting a batch to give a piece of it on each process, or by + yielding different full batches on each process. + + On two processes with a sampler of `[[0, 1, 2, 3], [4, 5, 6, 7]]`, this will result in: + + - the sampler on process 0 to yield `[0, 1, 2, 3]` and the sampler on process 1 to yield `[4, 5, 6, 7]` if + this argument is set to `False`. + - the sampler on process 0 to yield `[0, 1]` then `[4, 5]` and the sampler on process 1 to yield `[2, 3]` + then `[6, 7]` if this argument is set to `True`. + even_batches (`bool`, *optional*, defaults to `True`): + Whether or not to loop back at the beginning of the sampler when the number of samples is not a round + multiple of (original batch size / number of processes). + + + + `BatchSampler`s with varying batch sizes are not enabled by default. To enable this behaviour, set `even_batches` + equal to `False` + + """ + + def __init__( + self, + batch_sampler: BatchSampler, + num_processes: int = 1, + process_index: int = 0, + split_batches: bool = False, + even_batches: bool = True, + ): + if split_batches and batch_sampler.batch_size % num_processes != 0: + raise ValueError( + f"To use `BatchSamplerShard` in `split_batches` mode, the batch size ({batch_sampler.batch_size}) " + f"needs to be a round multiple of the number of processes ({num_processes})." + ) + self.batch_sampler = batch_sampler + self.num_processes = num_processes + self.process_index = process_index + self.split_batches = split_batches + self.even_batches = even_batches + self.batch_size = getattr(batch_sampler, "batch_size", None) + self.drop_last = getattr(batch_sampler, "drop_last", False) + if self.batch_size is None and self.even_batches: + raise ValueError( + "You need to use `even_batches=False` when the batch sampler has no batch size. If you " + "are not calling this method directly, set `accelerator.even_batches=False` instead." + ) + + @property + def total_length(self): + return len(self.batch_sampler) + + def __len__(self): + if self.split_batches: + # Split batches does not change the length of the batch sampler + return len(self.batch_sampler) + if len(self.batch_sampler) % self.num_processes == 0: + # If the length is a round multiple of the number of processes, it's easy. + return len(self.batch_sampler) // self.num_processes + length = len(self.batch_sampler) // self.num_processes + if self.drop_last: + # Same if we drop the remainder. + return length + elif self.even_batches: + # When we even batches we always get +1 + return length + 1 + else: + # Otherwise it depends on the process index. + return length + 1 if self.process_index < len(self.batch_sampler) % self.num_processes else length + + def __iter__(self): + return self._iter_with_split() if self.split_batches else self._iter_with_no_split() + + def _iter_with_split(self): + initial_data = [] + batch_length = self.batch_sampler.batch_size // self.num_processes + for idx, batch in enumerate(self.batch_sampler): + if idx == 0: + initial_data = batch + if len(batch) == self.batch_size: + # If the batch is full, we yield the part of it this process is responsible of. + yield batch[batch_length * self.process_index : batch_length * (self.process_index + 1)] + + # If drop_last is True of the last batch was full, iteration is over, otherwise... + if not self.drop_last and len(initial_data) > 0 and len(batch) < self.batch_size: + if not self.even_batches: + if len(batch) > batch_length * self.process_index: + yield batch[batch_length * self.process_index : batch_length * (self.process_index + 1)] + else: + # For degenerate cases where the dataset has less than num_process * batch_size samples + while len(initial_data) < self.batch_size: + initial_data += initial_data + batch = batch + initial_data + yield batch[batch_length * self.process_index : batch_length * (self.process_index + 1)] + + def _iter_with_no_split(self): + initial_data = [] + batch_to_yield = [] + for idx, batch in enumerate(self.batch_sampler): + # We gather the initial indices in case we need to circle back at the end. + if not self.drop_last and idx < self.num_processes: + initial_data += batch + # We identify the batch to yield but wait until we ar sure every process gets a full batch before actually + # yielding it. + if idx % self.num_processes == self.process_index: + batch_to_yield = batch + if idx % self.num_processes == self.num_processes - 1 and ( + self.batch_size is None or len(batch) == self.batch_size + ): + yield batch_to_yield + batch_to_yield = [] + + # If drop_last is True, iteration is over, otherwise... + if not self.drop_last and len(initial_data) > 0: + if not self.even_batches: + if len(batch_to_yield) > 0: + yield batch_to_yield + else: + # ... we yield the complete batch we had saved before if it has the proper length + if len(batch_to_yield) == self.batch_size: + yield batch_to_yield + + # For degenerate cases where the dataset has less than num_process * batch_size samples + while len(initial_data) < self.num_processes * self.batch_size: + initial_data += initial_data + + # If the last batch seen was of the proper size, it has been yielded by its process so we move to the next + if len(batch) == self.batch_size: + batch = [] + idx += 1 + + # Make sure we yield a multiple of self.num_processes batches + cycle_index = 0 + while idx % self.num_processes != 0 or len(batch) > 0: + end_index = cycle_index + self.batch_size - len(batch) + batch += initial_data[cycle_index:end_index] + if idx % self.num_processes == self.process_index: + yield batch + cycle_index = end_index + batch = [] + idx += 1 + + +class IterableDatasetShard(IterableDataset): + """ + Wraps a PyTorch `IterableDataset` to generate samples for one of the processes only. Instances of this class will + always yield a number of samples that is a round multiple of the actual batch size (depending of the value of + `split_batches`, this is either `batch_size` or `batch_size x num_processes`). Depending on the value of the + `drop_last` attribute of the batch sampler passed, it will either stop the iteration at the first batch that would + be too small or loop with indices from the beginning. + + Args: + dataset (`torch.utils.data.dataset.IterableDataset`): + The batch sampler to split in several shards. + batch_size (`int`, *optional*, defaults to 1): + The size of the batches per shard (if `split_batches=False`) or the size of the batches (if + `split_batches=True`). + drop_last (`bool`, *optional*, defaults to `False`): + Whether or not to drop the last incomplete batch or complete the last batches by using the samples from the + beginning. + num_processes (`int`, *optional*, defaults to 1): + The number of processes running concurrently. + process_index (`int`, *optional*, defaults to 0): + The index of the current process. + split_batches (`bool`, *optional*, defaults to `False`): + Whether the shards should be created by splitting a batch to give a piece of it on each process, or by + yielding different full batches on each process. + + On two processes with an iterable dataset yielding of `[0, 1, 2, 3, 4, 5, 6, 7]`, this will result in: + + - the shard on process 0 to yield `[0, 1, 2, 3]` and the shard on process 1 to yield `[4, 5, 6, 7]` if this + argument is set to `False`. + - the shard on process 0 to yield `[0, 1, 4, 5]` and the sampler on process 1 to yield `[2, 3, 6, 7]` if + this argument is set to `True`. + """ + + def __init__( + self, + dataset: IterableDataset, + batch_size: int = 1, + drop_last: bool = False, + num_processes: int = 1, + process_index: int = 0, + split_batches: bool = False, + ): + if split_batches and batch_size > 1 and batch_size % num_processes != 0: + raise ValueError( + f"To use `IterableDatasetShard` in `split_batches` mode, the batch size ({batch_size}) " + f"needs to be a round multiple of the number of processes ({num_processes})." + ) + self.dataset: IterableDataset = dataset + self.batch_size = batch_size + self.drop_last = drop_last + self.num_processes = num_processes + self.process_index = process_index + self.split_batches = split_batches + + def set_epoch(self, epoch): + self.epoch = epoch + if hasattr(self.dataset, "set_epoch"): + self.dataset.set_epoch(epoch) + + def __len__(self): + # We will just raise the downstream error if the underlying dataset is not sized + if self.drop_last: + return (len(self.dataset) // (self.batch_size * self.num_processes)) * self.batch_size + else: + return math.ceil(len(self.dataset) / (self.batch_size * self.num_processes)) * self.batch_size + + def __iter__(self): + if ( + not hasattr(self.dataset, "set_epoch") + and hasattr(self.dataset, "generator") + and isinstance(self.dataset.generator, torch.Generator) + ): + self.dataset.generator.manual_seed(self.epoch) + real_batch_size = self.batch_size if self.split_batches else (self.batch_size * self.num_processes) + process_batch_size = (self.batch_size // self.num_processes) if self.split_batches else self.batch_size + process_slice = range(self.process_index * process_batch_size, (self.process_index + 1) * process_batch_size) + + first_batch = None + current_batch = [] + for element in self.dataset: + current_batch.append(element) + # Wait to have a full batch before yielding elements. + if len(current_batch) == real_batch_size: + for i in process_slice: + yield current_batch[i] + if first_batch is None: + first_batch = current_batch.copy() + current_batch = [] + + # Finished if drop_last is True, otherwise complete the last batch with elements from the beginning. + if not self.drop_last and len(current_batch) > 0: + if first_batch is None: + first_batch = current_batch.copy() + while len(current_batch) < real_batch_size: + current_batch += first_batch + for i in process_slice: + yield current_batch[i] + + +class DataLoaderStateMixin: + """ + Mixin class that adds a state to a `DataLoader` to keep track of the status inside the dataloader such as at the + end of the iteration, the number of items in the dataset in the last batch relative to the batch size, and other + useful information that might be needed. + + **Available attributes:** + + - **end_of_dataloader** (`bool`) -- Whether at the last iteration or batch + - **remainder** (`int`) -- The number of items that are remaining in the last batch, relative to the total + batch size + + + + Inheriters of this class should ensure that the class creates a `GradientState()` instance, stored in + `self.gradient_state`. + + + + """ + + def __init_subclass__(cls, **kwargs): + cls.end_of_dataloader = False + cls.remainder = -1 + + def reset(self): + self.end_of_dataloader = False + self.remainder = -1 + + def begin(self): + "Prepares the gradient state for the current dataloader" + self.reset() + with suppress(Exception): + if not self._drop_last: + length = getattr(self.dataset, "total_dataset_length", len(self.dataset)) + self.remainder = length % self.total_batch_size + self.gradient_state._add_dataloader(self) + + def end(self): + "Cleans up the gradient state after exiting the dataloader" + self.gradient_state._remove_dataloader(self) + + +class DataLoaderAdapter: + """ + A class which wraps around a PyTorch `DataLoader` (or variants of it) to be used with the `Accelerator`. For + compatibility reasons, this class inherits from the class it wraps around, so it can be used as a drop-in. + """ + + def __init__(self, dataset, use_stateful_dataloader=False, batch_sampler=None, **kwargs): + self.use_stateful_dataloader = use_stateful_dataloader + if is_torchdata_stateful_dataloader_available(): + from torchdata.stateful_dataloader import StatefulDataLoader + + if use_stateful_dataloader and not is_torchdata_stateful_dataloader_available(): + raise ImportError( + "StatefulDataLoader is not available. Please install torchdata version 0.8.0 or higher to use it." + ) + if use_stateful_dataloader: + torchdata_version = version.parse(importlib.metadata.version("torchdata")) + if ( + "in_order" in kwargs + and compare_versions(torchdata_version, "<", "0.11") + and is_torch_version(">=", "2.6.0") + ): + kwargs.pop("in_order") + self.base_dataloader = StatefulDataLoader(dataset, batch_sampler=batch_sampler, **kwargs) + else: + self.base_dataloader = DataLoader(dataset, batch_sampler=batch_sampler, **kwargs) + + if hasattr(self.base_dataloader, "state_dict"): + self.dl_state_dict = self.base_dataloader.state_dict() + + def __getattr__(self, name): + # Avoid infinite recursion if we try to access a nonexistent base_dataloader attribute. + if name == "base_dataloader": + raise AttributeError() + # Delegate attribute access to the internal dataloader + return getattr(self.base_dataloader, name) + + def state_dict(self): + return self.dl_state_dict + + def load_state_dict(self, state_dict): + self.base_dataloader.load_state_dict(state_dict) + + @property + def __class__(self): + """ + In order to maintain backwards compatibility with other code, we need to ensure `isinstance(obj, DataLoader)` + returns true. This is because some downstream code assumes that the `DataLoader` is the base class of the + object. + """ + return self.base_dataloader.__class__ + + def __len__(self): + return len(self.base_dataloader) + + def adjust_state_dict_for_prefetch(self): + """ + Adjusts the state dict for prefetching. Natively, this will adjust all of the iters yielded keys in + `self.dl_state_dict` by a factor of `num_processes - 1`, however if a custom correction is needed, this can be + overridden. + + This should modify `self.dl_state_dict` directly + """ + # The state dict will be off by a factor of `n-1` batch too many during DDP, + # so we need to adjust it here + if PartialState().distributed_type != DistributedType.NO: + factor = PartialState().num_processes - 1 + if self.dl_state_dict["_sampler_iter_yielded"] > 0: + self.dl_state_dict["_sampler_iter_yielded"] -= factor + if self.dl_state_dict["_num_yielded"] > 0: + self.dl_state_dict["_num_yielded"] -= factor + if self.dl_state_dict["_index_sampler_state"] is not None: + if ( + "samples_yielded" in self.dl_state_dict["_index_sampler_state"] + and self.dl_state_dict["_index_sampler_state"]["samples_yielded"] > 0 + ): + self.dl_state_dict["_index_sampler_state"]["samples_yielded"] -= self.batch_size * factor + + def _update_state_dict(self): + # The state_dict of the underlying base_dataloader may be ahead of what is currently being yielded. + # E.g. the implementation of DataLoaderShard involves having an underlying iterator 1 element ahead of + # what it wants to yield. + # + # _update_state_dict is called to snapshot the state_dict that would properly recover the DataLoaderAdapter. + if hasattr(self.base_dataloader, "state_dict"): + self.dl_state_dict = self.base_dataloader.state_dict() + # Potentially modify the state_dict to adjust for prefetching + self.adjust_state_dict_for_prefetch() + # Then tag if we are at the end of the dataloader + self.dl_state_dict["_iterator_finished"] = self.end_of_dataloader + + +class DataLoaderShard(DataLoaderAdapter, DataLoaderStateMixin): + """ + Subclass of `DataLoaderAdapter` that will deal with device placement and current distributed setup. + + Args: + dataset (`torch.utils.data.dataset.Dataset`): + The dataset to use to build this dataloader. + device (`torch.device`, *optional*): + If passed, the device to put all batches on. + rng_types (list of `str` or [`~utils.RNGType`]): + The list of random number generators to synchronize at the beginning of each iteration. Should be one or + several of: + + - `"torch"`: the base torch random number generator + - `"cuda"`: the CUDA random number generator (GPU only) + - `"xla"`: the XLA random number generator (TPU only) + - `"generator"`: an optional `torch.Generator` + synchronized_generator (`torch.Generator`, *optional*): + A random number generator to keep synchronized across processes. + skip_batches (`int`, *optional*, defaults to 0): + The number of batches to skip at the beginning. + use_stateful_dataloader (`bool`, *optional*, defaults to `False`): + Whether to have this class adapt `StatefulDataLoader` from `torchdata` instead of the regular `DataLoader`. + **kwargs (additional keyword arguments, *optional*): + All other keyword arguments to pass to the regular `DataLoader` initialization. + + **Available attributes:** + + - **total_batch_size** (`int`) -- Total batch size of the dataloader across all processes. + Equal to the original batch size when `split_batches=True`; otherwise the original batch size * the total + number of processes + + - **total_dataset_length** (`int`) -- Total length of the inner dataset across all processes. + """ + + def __init__( + self, + dataset, + device=None, + rng_types=None, + synchronized_generator=None, + skip_batches=0, + use_stateful_dataloader=False, + _drop_last: bool = False, + _non_blocking: bool = False, + torch_device_mesh=None, + **kwargs, + ): + super().__init__(dataset, use_stateful_dataloader=use_stateful_dataloader, **kwargs) + self.device = device + self.rng_types = rng_types + self.synchronized_generator = synchronized_generator + self.skip_batches = skip_batches + self.gradient_state = GradientState() + self._drop_last = _drop_last + self._non_blocking = _non_blocking + self.iteration = 0 + + def __iter__(self): + if self.rng_types is not None: + synchronize_rng_states(self.rng_types, self.synchronized_generator) + self.begin() + + self.set_epoch(self.iteration) + dataloader_iter = self.base_dataloader.__iter__() + # We iterate one batch ahead to check when we are at the end + try: + current_batch = next(dataloader_iter) + except StopIteration: + self.end() + return + + batch_index = 0 + while True: + try: + # But we still move it to the device so it is done before `StopIteration` is reached + if self.device is not None: + current_batch = send_to_device(current_batch, self.device, non_blocking=self._non_blocking) + self._update_state_dict() + next_batch = next(dataloader_iter) + if batch_index >= self.skip_batches: + yield current_batch + batch_index += 1 + current_batch = next_batch + except StopIteration: + self.end_of_dataloader = True + self._update_state_dict() + if batch_index >= self.skip_batches: + yield current_batch + break + + self.iteration += 1 + self.end() + + def __reduce__(self): + """ + Define the `__reduce__` method to ensure a `DataLoaderShard` can be pickled and unpickled. This needs to be + explicitly defined since default pickling behavior is broken by `DataLoaderAdapter` messing with its + `__class__` member. + """ + args = super().__reduce__() + return (DataLoaderShard, *args[1:]) + + def set_epoch(self, epoch: int): + # In case it is manually passed in, the user can set it to what they like + if self.iteration != epoch: + self.iteration = epoch + if hasattr(self.batch_sampler, "set_epoch"): + self.batch_sampler.set_epoch(epoch) + if hasattr(self.batch_sampler, "sampler") and hasattr(self.batch_sampler.sampler, "set_epoch"): + self.batch_sampler.sampler.set_epoch(epoch) + if ( + hasattr(self.batch_sampler, "batch_sampler") + and hasattr(self.batch_sampler.batch_sampler, "sampler") + and hasattr(self.batch_sampler.batch_sampler.sampler, "set_epoch") + ): + self.batch_sampler.batch_sampler.sampler.set_epoch(epoch) + # We support if a custom `Dataset` implementation has `set_epoch` + # or in general HF datasets `Datasets` + elif hasattr(self.dataset, "set_epoch"): + self.dataset.set_epoch(epoch) + + @property + def total_batch_size(self): + batch_sampler = self.sampler if isinstance(self.sampler, BatchSampler) else self.batch_sampler + return ( + batch_sampler.batch_size + if getattr(batch_sampler, "split_batches", False) + else (batch_sampler.batch_size * getattr(batch_sampler, "num_processes", 1)) + ) + + @property + def total_dataset_length(self): + if hasattr(self.dataset, "total_length"): + return self.dataset.total_length + else: + return len(self.dataset) + + def get_sampler(self): + return get_sampler(self) + + def set_sampler(self, sampler): + sampler_is_batch_sampler = isinstance(self.sampler, BatchSampler) + if sampler_is_batch_sampler: + self.sampler.sampler = sampler + else: + self.batch_sampler.sampler = sampler + if hasattr(self.batch_sampler, "batch_sampler"): + self.batch_sampler.batch_sampler.sampler = sampler + + +if is_torch_xla_available(): + import torch_xla.distributed.parallel_loader as xpl + + class MpDeviceLoaderWrapper(xpl.MpDeviceLoader): + """ + Wrapper for the xpl.MpDeviceLoader class that knows the total batch size. + + XLA preloading threads will all call DataLoaderShard's __iter__(). Remove rng_types from DataLoaderShard to + prevent it from using the XLA device in the preloading threads, and synchronize the RNG once from the main + thread only. + + **Available attributes:** + + - **total_batch_size** (`int`) -- Total batch size of the dataloader across all processes. + Equal to the original batch size when `split_batches=True`; otherwise the original batch size * the total + number of processes + + - **total_dataset_length** (`int`) -- Total length of the inner dataset across all processes. + """ + + def __init__(self, dataloader: DataLoaderShard, device: torch.device): + super().__init__(dataloader, device) + self._rng_types = self._loader.rng_types + self._loader.rng_types = None + self.device = device + + def __iter__(self): + if self._rng_types is not None: + synchronize_rng_states(self._rng_types, self._loader.synchronized_generator) + + return super().__iter__() + + def set_epoch(self, epoch: int): + if hasattr(self.dataloader, "set_epoch"): + self.dataloader.set_epoch(epoch) + + @property + def total_batch_size(self): + return self._loader.total_batch_size + + @property + def total_dataset_length(self): + return self._loader.total_dataset_length + + @property + def batch_sampler(self): + return self._loader.batch_sampler + + @property + def dataloader(self): + return self._loader + + +class DataLoaderDispatcher(DataLoaderAdapter, DataLoaderStateMixin): + """ + Subclass of `DataLoaderAdapter` that will iterate and preprocess on process 0 only, then dispatch on each process + their part of the batch. + + Args: + split_batches (`bool`, *optional*, defaults to `False`): + Whether the resulting `DataLoader` should split the batches of the original data loader across devices or + yield full batches (in which case it will yield batches starting at the `process_index`-th and advancing of + `num_processes` batches at each iteration). Another way to see this is that the observed batch size will be + the same as the initial `dataloader` if this option is set to `True`, the batch size of the initial + `dataloader` multiplied by `num_processes` otherwise. Setting this option to `True` requires that the batch + size of the `dataloader` is a round multiple of `batch_size`. + skip_batches (`int`, *optional*, defaults to 0): + The number of batches to skip at the beginning of an iteration. + use_stateful_dataloader (`bool`, *optional*, defaults to `False`): + Whether to have this class adapt `StatefulDataLoader` from `torchdata` instead of the regular `DataLoader`. + + **Available attributes:** + + - **total_batch_size** (`int`) -- Total batch size of the dataloader across all processes. + Equal to the original batch size when `split_batches=True`; otherwise the original batch size * the total + number of processes + + - **total_dataset_length** (`int`) -- Total length of the inner dataset across all processes. + """ + + def __init__( + self, + dataset, + split_batches: bool = False, + skip_batches=0, + use_stateful_dataloader=False, + _drop_last: bool = False, + _non_blocking: bool = False, + slice_fn=None, + torch_device_mesh=None, + **kwargs, + ): + shuffle = False + from torch.utils.data.datapipes.iter.combinatorics import ShufflerIterDataPipe + + # We need to save the shuffling state of the DataPipe + if isinstance(dataset, ShufflerIterDataPipe): + shuffle = dataset._shuffle_enabled + super().__init__(dataset, use_stateful_dataloader=use_stateful_dataloader, **kwargs) + self.split_batches = split_batches + if shuffle: + torch.utils.data.graph_settings.apply_shuffle_settings(dataset, shuffle=shuffle) + + self.gradient_state = GradientState() + self.state = PartialState() + self._drop_last = _drop_last + self._non_blocking = _non_blocking + self.skip_batches = skip_batches + self.torch_device_mesh = torch_device_mesh + + self.slice_fn = slice_tensors if slice_fn is None else slice_fn + self.iteration = 0 + + # if a device mesh is provided extract each dimension (dp, fsdp, tp) + # device mesh may hold any number of dimensions, however, + # below code is for targeted support for dp, fsdp and tp + + # device mesh will be used only if there is tp involved + # or any multi-dimensional parallelism involving tp + # (dp, tp) (fsdp, tp) (dp, fsdp, tp) + # otherwise the default behaviour not using device mesh should be sufficient + # since multi dimensional parallelism devoid of tp would anyway need + # different batches for each process irrespective of dp or fsdp + self.submesh_tp = None + self.submesh_dp = None + self.submesh_fsdp = None + if self.torch_device_mesh and "tp" in self.torch_device_mesh.mesh_dim_names: + self.submesh_tp = self.torch_device_mesh["tp"] + if "dp" in self.torch_device_mesh.mesh_dim_names: + self.submesh_dp = self.torch_device_mesh["dp"] + if "fsdp" in self.torch_device_mesh.mesh_dim_names: + self.submesh_fsdp = self.torch_device_mesh["fsdp"] + if self.submesh_tp and (self.submesh_dp or self.submesh_fsdp): + raise ValueError("TP + (DP/FSDP) is not yet supported in dispatch mode") + + def _fetch_batches(self, iterator): + batches, batch = None, None + # On process 0, we gather the batch to dispatch. + if self.state.process_index == 0: + # Procedure to support TP only is simpler + # since we want to dispatch the same batch of samples across all ranks + # this removes complexity of handling multiple tp rank groups when TP + DP + # combination is involved. + + try: + # for TP case avoid using split_batches + # since it would mean that the dataloader should be spilling out + # duplicates of batches. + if self.split_batches: + # One batch of the main iterator is dispatched and split. + if self.submesh_tp: + logger.warning( + "Use of split_batches for TP would need the dataloader to produce duplicate batches," + "otherwise, use dispatch_batches=True instead." + ) + self._update_state_dict() + batch = next(iterator) + else: + # num_processes batches of the main iterator are concatenated then dispatched and split. + # We add the batches one by one so we have the remainder available when drop_last=False. + batches = [] + if self.submesh_tp: + # when tp, extract single batch and then replicate + self._update_state_dict() + batch = next(iterator) + batches = [batch] * self.state.num_processes + else: + for _ in range(self.state.num_processes): + self._update_state_dict() + batches.append(next(iterator)) + try: + batch = concatenate(batches, dim=0) + except RuntimeError as e: + raise RuntimeError( + "You can't use batches of different size with `dispatch_batches=True` or when using an `IterableDataset`." + "either pass `dispatch_batches=False` and have each process fetch its own batch " + " or pass `split_batches=True`. By doing so, the main process will fetch a full batch and " + "slice it into `num_processes` batches for each process." + ) from e + # In both cases, we need to get the structure of the batch that we will broadcast on other + # processes to initialize the tensors with the right shape. + # data_structure, stop_iteration + batch_info = [get_data_structure(batch), False] + except StopIteration: + batch_info = [None, True] + else: + batch_info = [None, self._stop_iteration] + # This is inplace, so after this instruction, every process has the same `batch_info` as process 0. + broadcast_object_list(batch_info) + self._stop_iteration = batch_info[1] + if self._stop_iteration: + # If drop_last is False and split_batches is False, we may have a remainder to take care of. + if not self.split_batches and not self._drop_last: + if self.state.process_index == 0 and len(batches) > 0: + batch = concatenate(batches, dim=0) + batch_info = [get_data_structure(batch), False] + else: + batch_info = [None, True] + broadcast_object_list(batch_info) + return batch, batch_info + + def __iter__(self): + self.begin() + self.set_epoch(self.iteration) + main_iterator = None + if is_torch_version(">=", "2.0.1"): + # NOTE PyTorch DataLoader adds forward compatibilities for DataPipes, which broadcasts + # shared seed to all dist processes. Thus, we need to create iterator for all dist processes. + # But, we only iterate through the DataLoader on process 0. + main_iterator = self.base_dataloader.__iter__() + elif self.state.process_index == 0: + main_iterator = self.base_dataloader.__iter__() + stop_iteration = False + self._stop_iteration = False + first_batch = None + next_batch, next_batch_info = self._fetch_batches(main_iterator) + batch_index = 0 + while not stop_iteration: + batch, batch_info = next_batch, next_batch_info + + if self.state.process_index != 0: + # Initialize tensors on other processes than process 0. + batch = initialize_tensors(batch_info[0]) + batch = send_to_device(batch, self.state.device, non_blocking=self._non_blocking) + # Broadcast the batch before splitting it. + batch = broadcast(batch, from_process=0) + + if not self._drop_last and first_batch is None: + # We keep at least num processes elements of the first batch to be able to complete the last batch + first_batch = self.slice_fn( + batch, + slice(0, self.state.num_processes), + process_index=self.state.process_index, + num_processes=self.state.num_processes, + ) + + if batch is None: + raise ValueError( + f"Batch does not contain any data (`{batch}`). At the end of all iterable data available before expected stop iteration." + ) + + observed_batch_size = find_batch_size(batch) + batch_size = observed_batch_size // self.state.num_processes + + stop_iteration = self._stop_iteration + if not stop_iteration: + # We may still be at the end of the dataloader without knowing it yet: if there is nothing left in + # the dataloader since the number of batches is a round multiple of the number of processes. + next_batch, next_batch_info = self._fetch_batches(main_iterator) + # next_batch_info[0] is None when there are no more batches, otherwise we still need to process them. + if self._stop_iteration and next_batch_info[0] is None: + stop_iteration = True + + if not self._drop_last and stop_iteration and observed_batch_size % self.state.num_processes != 0: + # If the last batch is not complete, let's add the first batch to it. + batch = concatenate([batch, first_batch], dim=0) + # Batch size computation above is wrong, it's off by 1 so we fix it. + batch_size += 1 + + data_slice = slice(self.state.process_index * batch_size, (self.state.process_index + 1) * batch_size) + batch = self.slice_fn( + batch, + data_slice, + process_index=self.state.process_index, + num_processes=self.state.num_processes, + ) + + if stop_iteration: + self.end_of_dataloader = True + self._update_state_dict() + self.remainder = observed_batch_size + if batch_index >= self.skip_batches: + yield batch + batch_index += 1 + self.iteration += 1 + self.end() + + def set_epoch(self, epoch: int): + # In case it is manually passed in, the user can set it to what they like + if self.iteration != epoch: + self.iteration = epoch + if hasattr(self.batch_sampler, "sampler") and hasattr(self.batch_sampler.sampler, "set_epoch"): + self.batch_sampler.sampler.set_epoch(epoch) + elif hasattr(self.dataset, "set_epoch"): + self.dataset.set_epoch(epoch) + + def __len__(self): + whole_length = len(self.base_dataloader) + if self.split_batches: + return whole_length + elif self._drop_last: + return whole_length // self.state.num_processes + else: + return math.ceil(whole_length / self.state.num_processes) + + def __reduce__(self): + """ + Define the `__reduce__` method to ensure a `DataLoaderDispatcher` can be pickled and unpickled. This needs to + be explicitly defined since default pickling behavior is broken by `DataLoaderAdapter` messing with its + `__class__` member. + """ + args = super().__reduce__() + return (DataLoaderDispatcher, *args[1:]) + + @property + def total_batch_size(self): + return ( + self.dataset.batch_size if self.split_batches else (self.dataset.batch_size * self.dataset.num_processes) + ) + + @property + def total_dataset_length(self): + return len(self.dataset) + + def get_sampler(self): + return get_sampler(self) + + def set_sampler(self, sampler): + sampler_is_batch_sampler = isinstance(self.sampler, BatchSampler) + if sampler_is_batch_sampler: + self.sampler.sampler = sampler + else: + self.batch_sampler.sampler = sampler + if hasattr(self.batch_sampler, "batch_sampler"): + self.batch_sampler.batch_sampler.sampler = sampler + + +def get_sampler(dataloader): + """ + Get the sampler associated to the dataloader + + Args: + dataloader (`torch.utils.data.dataloader.DataLoader`): + The data loader to split across several devices. + Returns: + `torch.utils.data.Sampler`: The sampler associated to the dataloader + """ + sampler_is_batch_sampler = isinstance(dataloader.sampler, BatchSampler) + if sampler_is_batch_sampler: + sampler = getattr(dataloader.sampler, "sampler", None) + else: + sampler = getattr(dataloader.batch_sampler, "sampler", None) + return sampler + + +def prepare_data_loader( + dataloader: DataLoader, + device: Optional[torch.device] = None, + num_processes: Optional[int] = None, + process_index: Optional[int] = None, + split_batches: bool = False, + put_on_device: bool = False, + rng_types: Optional[list[Union[str, RNGType]]] = None, + dispatch_batches: Optional[bool] = None, + even_batches: bool = True, + slice_fn_for_dispatch: Optional[Callable] = None, + use_seedable_sampler: bool = False, + data_seed: Optional[int] = None, + non_blocking: bool = False, + use_stateful_dataloader: bool = False, + torch_device_mesh=None, +) -> DataLoader: + """ + Wraps a PyTorch `DataLoader` to generate batches for one of the processes only. + + Depending on the value of the `drop_last` attribute of the `dataloader` passed, it will either stop the iteration + at the first batch that would be too small / not present on all processes or loop with indices from the beginning. + + Args: + dataloader (`torch.utils.data.dataloader.DataLoader`): + The data loader to split across several devices. + device (`torch.device`): + The target device for the returned `DataLoader`. + num_processes (`int`, *optional*): + The number of processes running concurrently. Will default to the value given by [`~state.PartialState`]. + process_index (`int`, *optional*): + The index of the current process. Will default to the value given by [`~state.PartialState`]. + split_batches (`bool`, *optional*, defaults to `False`): + Whether the resulting `DataLoader` should split the batches of the original data loader across devices or + yield full batches (in which case it will yield batches starting at the `process_index`-th and advancing of + `num_processes` batches at each iteration). + + Another way to see this is that the observed batch size will be the same as the initial `dataloader` if + this option is set to `True`, the batch size of the initial `dataloader` multiplied by `num_processes` + otherwise. + + Setting this option to `True` requires that the batch size of the `dataloader` is a round multiple of + `batch_size`. + put_on_device (`bool`, *optional*, defaults to `False`): + Whether or not to put the batches on `device` (only works if the batches are nested list, tuples or + dictionaries of tensors). + rng_types (list of `str` or [`~utils.RNGType`]): + The list of random number generators to synchronize at the beginning of each iteration. Should be one or + several of: + + - `"torch"`: the base torch random number generator + - `"cuda"`: the CUDA random number generator (GPU only) + - `"xla"`: the XLA random number generator (TPU only) + - `"generator"`: the `torch.Generator` of the sampler (or batch sampler if there is no sampler in your + dataloader) or of the iterable dataset (if it exists) if the underlying dataset is of that type. + + dispatch_batches (`bool`, *optional*): + If set to `True`, the dataloader prepared is only iterated through on the main process and then the batches + are split and broadcast to each process. Will default to `True` when the underlying dataset is an + `IterableDataset`, `False` otherwise. + even_batches (`bool`, *optional*, defaults to `True`): + If set to `True`, in cases where the total batch size across all processes does not exactly divide the + dataset, samples at the start of the dataset will be duplicated so the batch can be divided equally among + all workers. + slice_fn_for_dispatch (`Callable`, *optional*`): + If passed, this function will be used to slice tensors across `num_processes`. Will default to + [`~utils.slice_tensors`]. This argument is used only when `dispatch_batches` is set to `True` and will be + ignored otherwise. + use_seedable_sampler (`bool`, *optional*, defaults to `False`): + Whether to use the [`~data_loader.SeedableRandomSampler`] instead of a `RandomSampler` for better + reproducibility. Comes at a cost of potentially different performances due to different shuffling + algorithms but ensures results will be the *exact* same. Should be paired with `set_seed()` at every + `self.set_epoch` + data_seed (`int`, *optional*, defaults to `None`): + The seed to use for the underlying generator when using `use_seedable_sampler`. If `None`, the generator + will use the current default seed from torch. + non_blocking (`bool`, *optional*, defaults to `False`): + If set to `True`, dataloader will utilize non-blocking host-to-device transfers. If the dataloader has + `pin_memory` set to `True`, this will help to increase overlap between data transfer and computations. + use_stateful_dataloader (`bool`, *optional*, defaults to `False`): + "If set to true, the dataloader prepared by the Accelerator will be backed by " + "[torchdata.StatefulDataLoader](https://github.com/pytorch/data/tree/main/torchdata/stateful_dataloader). + This requires `torchdata` version 0.8.0 or higher that supports StatefulDataLoader to be installed." + torch_device_mesh (`torch.distributed.DeviceMesh`, *optional*, defaults to `None`): + PyTorch device mesh. + + + Returns: + `torch.utils.data.dataloader.DataLoader`: A new data loader that will yield the portion of the batches + + + + `BatchSampler`s with varying batch sizes are not enabled by default. To enable this behaviour, set `even_batches` + equal to `False` + + + """ + if dispatch_batches is None: + if not put_on_device: + dispatch_batches = False + else: + dispatch_batches = isinstance(dataloader.dataset, IterableDataset) + + if dispatch_batches and not put_on_device: + raise ValueError("Using `dispatch_batches=True` requires `put_on_device=True`.") + # Grab defaults from PartialState + state = PartialState() + if num_processes is None: + num_processes = state.num_processes + + if process_index is None: + process_index = state.process_index + + if torch_device_mesh: + if state.distributed_type == DistributedType.DEEPSPEED: + # In DeepSpeed, the optimizer sharing level in DP is determined by the config file. + # Only considers "dp" and "tp". + # Given a device mesh (dp, tp) = (2, 3): + # - From the data parallel perspective, ranks should be structured as: 0 0 0 1 1 1 + # - Processes with the same DP rank will receive the same batch. + submesh_tp_size = 1 + if "tp" in torch_device_mesh.mesh_dim_names: + submesh_tp_size = torch_device_mesh["tp"].size() + process_index = process_index // submesh_tp_size + num_processes = num_processes // submesh_tp_size + else: + # when device mesh is used, specifically with TP + # then there is need to update process_index and num_processes + # to bring in the effect of generating same batch across TP ranks + # and different batch across FSDP and DP ranks. + # Example: + # if device mesh is (dp,fsdp,tp) = (2, 2, 3) + # ranks would range from 0...11 + # from data angle ranks should look like 0 0 0 1 1 1 2 2 2 3 3 3 + # processes with same ranks/ids would receive the same batch + # for CP the same as TP applies + submesh_fsdp_size = 1 + submesh_dp_size = 1 + submesh_tp_size = 1 + submesh_cp_size = 1 + if "tp" in torch_device_mesh.mesh_dim_names: + submesh_tp_size = torch_device_mesh["tp"].size() + if "cp" in torch_device_mesh.mesh_dim_names: + submesh_cp_size = torch_device_mesh["cp"].size() + if "dp_replicate" in torch_device_mesh.mesh_dim_names: + submesh_dp_size = torch_device_mesh["dp_replicate"].size() + if "dp_shard" in torch_device_mesh.mesh_dim_names: + submesh_fsdp_size = torch_device_mesh["dp_shard"].size() + process_index = process_index // (submesh_tp_size * submesh_cp_size) + num_processes = submesh_fsdp_size * submesh_dp_size + + # Sanity check + if split_batches: + if dataloader.batch_size is not None: + batch_size_for_check = dataloader.batch_size + else: + # For custom batch_sampler + if hasattr(dataloader.batch_sampler, "batch_size"): + batch_size_for_check = dataloader.batch_sampler.batch_size + else: + raise ValueError( + "In order to use `split_batches==True` you must have a `batch_size` attribute either in the passed " + "`dataloader` or `dataloader.batch_sampler` objects, and it has to return a natural number. " + "Your `dataloader.batch_size` is None and `dataloader.batch_sampler` " + f"(`{type(dataloader.batch_sampler)}`) does not have the `batch_size` attribute set." + ) + + if batch_size_for_check > 1 and batch_size_for_check % num_processes != 0: + raise ValueError( + f"To use a `DataLoader` in `split_batches` mode, the batch size ({dataloader.batch_size}) " + f"needs to be a round multiple of the number of processes ({num_processes})." + ) + + new_dataset = dataloader.dataset + # Iterable dataset doesn't like batch_sampler, but data_loader creates a default one for it + new_batch_sampler = dataloader.batch_sampler if not isinstance(new_dataset, IterableDataset) else None + sampler_is_batch_sampler = isinstance(dataloader.sampler, BatchSampler) + synchronized_generator = None + + sampler = get_sampler(dataloader) + if isinstance(sampler, RandomSampler) and use_seedable_sampler: + # When iterating through the dataloader during distributed processes + # we want to ensure that on each process we are iterating through the same + # samples in the same order if a seed is set. This requires a tweak + # to the `torch.utils.data.RandomSampler` class (if used). + sampler = SeedableRandomSampler( + data_source=sampler.data_source, + replacement=sampler.replacement, + num_samples=sampler._num_samples, + generator=getattr( + sampler, + "generator", + torch.Generator(device=torch.get_default_device() if hasattr(torch, "get_default_device") else "cpu"), + ), + data_seed=data_seed, + ) + + if isinstance(dataloader.sampler, RandomSampler) and state.distributed_type == DistributedType.XLA: + # isinstance(dataloader.sampler, RandomSampler) indicates the original dataloader has `shuffle` enabled. + generator = torch.Generator( + device=torch.get_default_device() if hasattr(torch, "get_default_device") else "cpu" + ) + seed = int(torch.empty((), dtype=torch.int64).random_().item()) + generator.manual_seed(seed) + dataloader.generator = generator + dataloader.sampler.generator = generator + # No change if no multiprocess + if (num_processes != 1 or state.distributed_type == DistributedType.MEGATRON_LM) and not dispatch_batches: + if is_datasets_available(): + from datasets import IterableDataset as DatasetsIterableDataset + if ( + is_datasets_available() + and isinstance(new_dataset, DatasetsIterableDataset) + and not split_batches + and new_dataset.n_shards > num_processes + ): + new_dataset = new_dataset.shard(num_shards=num_processes, index=process_index) + elif isinstance(new_dataset, IterableDataset): + if getattr(dataloader.dataset, "generator", None) is not None: + synchronized_generator = dataloader.dataset.generator + new_dataset = IterableDatasetShard( + new_dataset, + batch_size=dataloader.batch_size, + drop_last=dataloader.drop_last, + num_processes=num_processes, + process_index=process_index, + split_batches=split_batches, + ) + else: + if not use_seedable_sampler and hasattr(sampler, "generator"): + if sampler.generator is None: + sampler.generator = torch.Generator( + device=torch.get_default_device() if hasattr(torch, "get_default_device") else "cpu" + ) + seed = int(torch.empty((), dtype=torch.int64).random_().item()) + sampler.generator.manual_seed(seed) + synchronized_generator = sampler.generator + batch_sampler = dataloader.sampler if sampler_is_batch_sampler else dataloader.batch_sampler + new_batch_sampler = BatchSamplerShard( + batch_sampler, + num_processes=num_processes, + process_index=process_index, + split_batches=split_batches, + even_batches=even_batches, + ) + + # We ignore all of those since they are all dealt with by our new_batch_sampler + ignore_kwargs = [ + "batch_size", + "shuffle", + "sampler", + "batch_sampler", + "drop_last", + ] + + if rng_types is not None and synchronized_generator is None and "generator" in rng_types: + rng_types.remove("generator") + + kwargs = { + k: getattr(dataloader, k, _PYTORCH_DATALOADER_KWARGS[k]) + for k in _PYTORCH_DATALOADER_KWARGS + if k not in ignore_kwargs + } + + # Need to provide batch_size as batch_sampler is None for Iterable dataset + if new_batch_sampler is None: + kwargs["drop_last"] = dataloader.drop_last + kwargs["batch_size"] = ( + dataloader.batch_size // num_processes if split_batches and not dispatch_batches else dataloader.batch_size + ) + if dispatch_batches: + kwargs.pop("generator") + dataloader = DataLoaderDispatcher( + new_dataset, + split_batches=split_batches, + batch_sampler=new_batch_sampler, + _drop_last=dataloader.drop_last, + _non_blocking=non_blocking, + slice_fn=slice_fn_for_dispatch, + use_stateful_dataloader=use_stateful_dataloader, + torch_device_mesh=torch_device_mesh, + **kwargs, + ) + elif sampler_is_batch_sampler: + dataloader = DataLoaderShard( + new_dataset, + device=device if put_on_device and state.distributed_type != DistributedType.XLA else None, + sampler=new_batch_sampler, + batch_size=dataloader.batch_size, + rng_types=rng_types, + _drop_last=dataloader.drop_last, + _non_blocking=non_blocking, + synchronized_generator=synchronized_generator, + use_stateful_dataloader=use_stateful_dataloader, + **kwargs, + ) + else: + dataloader = DataLoaderShard( + new_dataset, + device=device if put_on_device and state.distributed_type != DistributedType.XLA else None, + batch_sampler=new_batch_sampler, + rng_types=rng_types, + synchronized_generator=synchronized_generator, + _drop_last=dataloader.drop_last, + _non_blocking=non_blocking, + use_stateful_dataloader=use_stateful_dataloader, + **kwargs, + ) + + if isinstance(sampler, SeedableRandomSampler) and use_seedable_sampler: + dataloader.set_sampler(sampler) + if state.distributed_type == DistributedType.XLA: + return MpDeviceLoaderWrapper(dataloader, device) + return dataloader + + +class SkipBatchSampler(BatchSampler): + """ + A `torch.utils.data.BatchSampler` that skips the first `n` batches of another `torch.utils.data.BatchSampler`. + Should not be used if the original dataloader is a `StatefulDataLoader`. + """ + + def __init__(self, batch_sampler, skip_batches=0): + self.batch_sampler = batch_sampler + self.skip_batches = skip_batches + + def __iter__(self): + for index, samples in enumerate(self.batch_sampler): + if index >= self.skip_batches: + yield samples + + @property + def total_length(self): + return len(self.batch_sampler) + + def __len__(self): + return len(self.batch_sampler) - self.skip_batches + + +class SkipDataLoader(DataLoaderAdapter, DataLoaderStateMixin): + """ + Subclass of a PyTorch `DataLoader` that will skip the first batches. Generally it's preferable to use + `skip_first_batches`/`torchdata.StatefulDataLoader` instead of this class. + + Args: + dataset (`torch.utils.data.dataset.Dataset`): + The dataset to use to build this dataloader. + skip_batches (`int`, *optional*, defaults to 0): + The number of batches to skip at the beginning. + kwargs: + All other keyword arguments to pass to the regular `DataLoader` initialization. + """ + + def __init__(self, dataset, skip_batches=0, use_stateful_dataloader=False, **kwargs): + super().__init__(dataset, use_stateful_dataloader=use_stateful_dataloader, **kwargs) + self.skip_batches = skip_batches + self.gradient_state = GradientState() + + def __iter__(self): + self.begin() + for index, batch in enumerate(self.base_dataloader.__iter__()): + if index >= self.skip_batches: + self._update_state_dict() + yield batch + self.end() + + def __len__(self): + return len(self.base_dataloader) - self.skip_batches + + def __reduce__(self): + """ + Define the `__reduce__` method to ensure a `SkipDataLoader` can be pickled and unpickled. This needs to be + explicitly defined since default pickling behavior is broken by `DataLoaderAdapter` messing with its + `__class__` member. + """ + args = super().__reduce__() + return (SkipDataLoader, *args[1:]) + + +def skip_first_batches(dataloader, num_batches=0): + """ + Creates a `torch.utils.data.DataLoader` that will efficiently skip the first `num_batches`. Should not be used if + the original dataloader is a `StatefulDataLoader`. + """ + state = PartialState() + if state.distributed_type == DistributedType.XLA: + device = dataloader.device + dataloader = dataloader.dataloader + + dataset = dataloader.dataset + sampler_is_batch_sampler = False + if isinstance(dataset, IterableDataset): + new_batch_sampler = None + else: + sampler_is_batch_sampler = isinstance(dataloader.sampler, BatchSampler) + batch_sampler = dataloader.sampler if sampler_is_batch_sampler else dataloader.batch_sampler + new_batch_sampler = SkipBatchSampler(batch_sampler, skip_batches=num_batches) + + # We ignore all of those since they are all dealt with by our new_batch_sampler + ignore_kwargs = [ + "batch_size", + "shuffle", + "sampler", + "batch_sampler", + "drop_last", + ] + + kwargs = { + k: getattr(dataloader, k, _PYTORCH_DATALOADER_KWARGS[k]) + for k in _PYTORCH_DATALOADER_KWARGS + if k not in ignore_kwargs + } + + # Need to provide batch_size as batch_sampler is None for Iterable dataset + if new_batch_sampler is None: + kwargs["drop_last"] = dataloader.drop_last + kwargs["batch_size"] = dataloader.batch_size + + if isinstance(dataloader, DataLoaderDispatcher): + if new_batch_sampler is None: + # Need to manually skip batches in the dataloader + kwargs["skip_batches"] = num_batches + dataloader = DataLoaderDispatcher( + dataset, + split_batches=dataloader.split_batches, + batch_sampler=new_batch_sampler, + _drop_last=dataloader._drop_last, + **kwargs, + ) + elif isinstance(dataloader, DataLoaderShard): + if new_batch_sampler is None: + # Need to manually skip batches in the dataloader + kwargs["skip_batches"] = num_batches + elif sampler_is_batch_sampler: + kwargs["sampler"] = new_batch_sampler + kwargs["batch_size"] = dataloader.batch_size + else: + kwargs["batch_sampler"] = new_batch_sampler + dataloader = DataLoaderShard( + dataset, + device=dataloader.device, + rng_types=dataloader.rng_types, + synchronized_generator=dataloader.synchronized_generator, + **kwargs, + ) + else: + if new_batch_sampler is None: + # Need to manually skip batches in the dataloader + dataloader = SkipDataLoader(dataset, skip_batches=num_batches, **kwargs) + else: + dataloader = DataLoader(dataset, batch_sampler=new_batch_sampler, **kwargs) + + if state.distributed_type == DistributedType.XLA: + dataloader = MpDeviceLoaderWrapper(dataloader, device) + + return dataloader diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/hooks.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/hooks.py new file mode 100644 index 0000000000000000000000000000000000000000..5d6a22e0a9e7db49edea04d6c330477ef981e8bc --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/hooks.py @@ -0,0 +1,776 @@ +# Copyright 2022 The HuggingFace Team. All rights reserved. +# +# 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. + +import functools +from collections.abc import Mapping +from typing import Optional, Union + +import torch +import torch.nn as nn + +from .state import PartialState +from .utils import ( + PrefixedDataset, + find_device, + named_module_tensors, + send_to_device, + set_module_tensor_to_device, +) +from .utils.imports import ( + is_mlu_available, + is_musa_available, + is_npu_available, +) +from .utils.memory import clear_device_cache +from .utils.modeling import get_non_persistent_buffers +from .utils.other import recursive_getattr + + +_accelerate_added_attributes = ["to", "cuda", "npu", "xpu", "mlu", "sdaa", "musa"] + + +class ModelHook: + """ + A hook that contains callbacks to be executed just before and after the forward method of a model. The difference + with PyTorch existing hooks is that they get passed along the kwargs. + + Class attribute: + - **no_grad** (`bool`, *optional*, defaults to `False`) -- Whether or not to execute the actual forward pass under + the `torch.no_grad()` context manager. + """ + + no_grad = False + + def init_hook(self, module): + """ + To be executed when the hook is attached to the module. + + Args: + module (`torch.nn.Module`): The module attached to this hook. + """ + return module + + def pre_forward(self, module, *args, **kwargs): + """ + To be executed just before the forward method of the model. + + Args: + module (`torch.nn.Module`): The module whose forward pass will be executed just after this event. + args (`Tuple[Any]`): The positional arguments passed to the module. + kwargs (`Dict[Str, Any]`): The keyword arguments passed to the module. + + Returns: + `Tuple[Tuple[Any], Dict[Str, Any]]`: A tuple with the treated `args` and `kwargs`. + """ + return args, kwargs + + def post_forward(self, module, output): + """ + To be executed just after the forward method of the model. + + Args: + module (`torch.nn.Module`): The module whose forward pass been executed just before this event. + output (`Any`): The output of the module. + + Returns: + `Any`: The processed `output`. + """ + return output + + def detach_hook(self, module): + """ + To be executed when the hook is detached from a module. + + Args: + module (`torch.nn.Module`): The module detached from this hook. + """ + return module + + +class SequentialHook(ModelHook): + """ + A hook that can contain several hooks and iterates through them at each event. + """ + + def __init__(self, *hooks): + self.hooks = hooks + + def init_hook(self, module): + for hook in self.hooks: + module = hook.init_hook(module) + return module + + def pre_forward(self, module, *args, **kwargs): + for hook in self.hooks: + args, kwargs = hook.pre_forward(module, *args, **kwargs) + return args, kwargs + + def post_forward(self, module, output): + for hook in self.hooks: + output = hook.post_forward(module, output) + return output + + def detach_hook(self, module): + for hook in self.hooks: + module = hook.detach_hook(module) + return module + + +def add_hook_to_module(module: nn.Module, hook: ModelHook, append: bool = False): + """ + Adds a hook to a given module. This will rewrite the `forward` method of the module to include the hook, to remove + this behavior and restore the original `forward` method, use `remove_hook_from_module`. + + + + If the module already contains a hook, this will replace it with the new hook passed by default. To chain two hooks + together, pass `append=True`, so it chains the current and new hook into an instance of the `SequentialHook` class. + + + + Args: + module (`torch.nn.Module`): + The module to attach a hook to. + hook (`ModelHook`): + The hook to attach. + append (`bool`, *optional*, defaults to `False`): + Whether the hook should be chained with an existing one (if module already contains a hook) or not. + + Returns: + `torch.nn.Module`: The same module, with the hook attached (the module is modified in place, so the result can + be discarded). + """ + if append and (getattr(module, "_hf_hook", None) is not None): + old_hook = module._hf_hook + remove_hook_from_module(module) + hook = SequentialHook(old_hook, hook) + + if hasattr(module, "_hf_hook") and hasattr(module, "_old_forward"): + # If we already put some hook on this module, we replace it with the new one. + old_forward = module._old_forward + else: + old_forward = module.forward + module._old_forward = old_forward + + module = hook.init_hook(module) + module._hf_hook = hook + + def new_forward(module, *args, **kwargs): + args, kwargs = module._hf_hook.pre_forward(module, *args, **kwargs) + if module._hf_hook.no_grad: + with torch.no_grad(): + output = module._old_forward(*args, **kwargs) + else: + output = module._old_forward(*args, **kwargs) + return module._hf_hook.post_forward(module, output) + + # Overriding a GraphModuleImpl forward freezes the forward call and later modifications on the graph will fail. + # Reference: https://pytorch.slack.com/archives/C3PDTEV8E/p1705929610405409 + if "GraphModuleImpl" in str(type(module)): + module.__class__.forward = functools.update_wrapper(functools.partial(new_forward, module), old_forward) + else: + module.forward = functools.update_wrapper(functools.partial(new_forward, module), old_forward) + + return module + + +def remove_hook_from_module(module: nn.Module, recurse=False): + """ + Removes any hook attached to a module via `add_hook_to_module`. + + Args: + module (`torch.nn.Module`): The module to attach a hook to. + recurse (`bool`, **optional**): Whether to remove the hooks recursively + + Returns: + `torch.nn.Module`: The same module, with the hook detached (the module is modified in place, so the result can + be discarded). + """ + + if hasattr(module, "_hf_hook"): + module._hf_hook.detach_hook(module) + delattr(module, "_hf_hook") + + if hasattr(module, "_old_forward"): + # Overriding a GraphModuleImpl forward freezes the forward call and later modifications on the graph will fail. + # Reference: https://pytorch.slack.com/archives/C3PDTEV8E/p1705929610405409 + if "GraphModuleImpl" in str(type(module)): + module.__class__.forward = module._old_forward + else: + module.forward = module._old_forward + delattr(module, "_old_forward") + + # Remove accelerate added warning hooks from dispatch_model + for attr in _accelerate_added_attributes: + module.__dict__.pop(attr, None) + + if recurse: + for child in module.children(): + remove_hook_from_module(child, recurse) + + return module + + +class AlignDevicesHook(ModelHook): + """ + A generic `ModelHook` that ensures inputs and model weights are on the same device for the forward pass of the + associated module, potentially offloading the weights after the forward pass. + + Args: + execution_device (`torch.device`, *optional*): + The device on which inputs and model weights should be placed before the forward pass. + offload (`bool`, *optional*, defaults to `False`): + Whether or not the weights should be offloaded after the forward pass. + io_same_device (`bool`, *optional*, defaults to `False`): + Whether or not the output should be placed on the same device as the input was. + weights_map (`Mapping[str, torch.Tensor]`, *optional*): + When the model weights are offloaded, a (potentially lazy) map from param names to the tensor values. + offload_buffers (`bool`, *optional*, defaults to `False`): + Whether or not to include the associated module's buffers when offloading. + place_submodules (`bool`, *optional*, defaults to `False`): + Whether to place the submodules on `execution_device` during the `init_hook` event. + """ + + def __init__( + self, + execution_device: Optional[Union[int, str, torch.device]] = None, + offload: bool = False, + io_same_device: bool = False, + weights_map: Optional[Mapping] = None, + offload_buffers: bool = False, + place_submodules: bool = False, + skip_keys: Optional[Union[str, list[str]]] = None, + tied_params_map: Optional[dict[int, dict[torch.device, torch.Tensor]]] = None, + ): + self.execution_device = execution_device + self.offload = offload + self.io_same_device = io_same_device + self.weights_map = weights_map + self.offload_buffers = offload_buffers + self.place_submodules = place_submodules + self.skip_keys = skip_keys + + # Will contain the input device when `io_same_device=True`. + self.input_device = None + self.param_original_devices = {} + self.buffer_original_devices = {} + self.tied_params_names = set() + + # The hook pre_forward/post_forward need to have knowledge of this dictionary, as with offloading we want to avoid duplicating memory + # for tied weights already loaded on the target execution device. + self.tied_params_map = tied_params_map + + def __repr__(self): + return ( + f"AlignDevicesHook(execution_device={self.execution_device}, offload={self.offload}, " + f"io_same_device={self.io_same_device}, offload_buffers={self.offload_buffers}, " + f"place_submodules={self.place_submodules}, skip_keys={repr(self.skip_keys)})" + ) + + def init_hook(self, module): + # In case the AlignDevicesHook is on meta device, ignore tied weights as data_ptr() is then always zero. + if self.execution_device == "meta" or self.execution_device == torch.device("meta"): + self.tied_params_map = None + + if not self.offload and self.execution_device is not None: + for name, _ in named_module_tensors(module, recurse=self.place_submodules): + set_module_tensor_to_device(module, name, self.execution_device, tied_params_map=self.tied_params_map) + elif self.offload: + self.original_devices = { + name: param.device for name, param in named_module_tensors(module, recurse=self.place_submodules) + } + if self.weights_map is None: + self.weights_map = { + name: param.to("cpu") + for name, param in named_module_tensors( + module, include_buffers=self.offload_buffers, recurse=self.place_submodules + ) + } + for name, _ in named_module_tensors( + module, include_buffers=self.offload_buffers, recurse=self.place_submodules, remove_non_persistent=True + ): + # When using disk offloading, we can not rely on `weights_map[name].data_ptr()` as the reference pointer, + # as we have no guarantee that safetensors' `file.get_tensor()` will always give the same pointer. + # As we have no reliable way to track the shared data pointer of tied weights in this case, we use tied_params_names: List[str] + # to add on the fly pointers to `tied_params_map` in the pre_forward call. + if ( + self.tied_params_map is not None + and recursive_getattr(module, name).data_ptr() in self.tied_params_map + ): + self.tied_params_names.add(name) + + set_module_tensor_to_device(module, name, "meta") + + if not self.offload_buffers and self.execution_device is not None: + for name, _ in module.named_buffers(recurse=self.place_submodules): + set_module_tensor_to_device( + module, name, self.execution_device, tied_params_map=self.tied_params_map + ) + elif self.offload_buffers and self.execution_device is not None: + for name in get_non_persistent_buffers(module, recurse=self.place_submodules): + set_module_tensor_to_device( + module, name, self.execution_device, tied_params_map=self.tied_params_map + ) + + return module + + def pre_forward(self, module, *args, **kwargs): + if self.io_same_device: + self.input_device = find_device([args, kwargs]) + if self.offload: + self.tied_pointers_to_remove = set() + + for name, _ in named_module_tensors( + module, + include_buffers=self.offload_buffers, + recurse=self.place_submodules, + remove_non_persistent=True, + ): + fp16_statistics = None + value = self.weights_map[name] + if "weight" in name and name.replace("weight", "SCB") in self.weights_map.keys(): + if value.dtype == torch.int8: + fp16_statistics = self.weights_map[name.replace("weight", "SCB")] + + # In case we are using offloading with tied weights, we need to keep track of the offloaded weights + # that are loaded on device at this point, as we will need to remove them as well from the dictionary + # self.tied_params_map in order to allow to free memory. + if name in self.tied_params_names and value.data_ptr() not in self.tied_params_map: + self.tied_params_map[value.data_ptr()] = {} + + if ( + value is not None + and self.tied_params_map is not None + and value.data_ptr() in self.tied_params_map + and self.execution_device not in self.tied_params_map[value.data_ptr()] + ): + self.tied_pointers_to_remove.add((value.data_ptr(), self.execution_device)) + + set_module_tensor_to_device( + module, + name, + self.execution_device, + value=value, + fp16_statistics=fp16_statistics, + tied_params_map=self.tied_params_map, + ) + + return send_to_device(args, self.execution_device), send_to_device( + kwargs, self.execution_device, skip_keys=self.skip_keys + ) + + def post_forward(self, module, output): + if self.offload: + for name, _ in named_module_tensors( + module, + include_buffers=self.offload_buffers, + recurse=self.place_submodules, + remove_non_persistent=True, + ): + set_module_tensor_to_device(module, name, "meta") + if type(module).__name__ == "Linear8bitLt": + module.state.SCB = None + module.state.CxB = None + + # We may have loaded tied weights into self.tied_params_map (avoiding to load them several times in e.g. submodules): remove them from + # this dictionary to allow the garbage collector to do its job. + for value_pointer, device in self.tied_pointers_to_remove: + if isinstance(device, int): + if is_npu_available(): + device = f"npu:{device}" + elif is_mlu_available(): + device = f"mlu:{device}" + elif is_musa_available(): + device = f"musa:{device}" + if device in self.tied_params_map[value_pointer]: + del self.tied_params_map[value_pointer][device] + self.tied_pointers_to_remove = set() + if self.io_same_device and self.input_device is not None: + output = send_to_device(output, self.input_device, skip_keys=self.skip_keys) + + return output + + def detach_hook(self, module): + if self.offload: + for name, device in self.original_devices.items(): + if device != torch.device("meta"): + set_module_tensor_to_device(module, name, device, value=self.weights_map.get(name, None)) + return module + + +def attach_execution_device_hook( + module: torch.nn.Module, + execution_device: Union[int, str, torch.device], + skip_keys: Optional[Union[str, list[str]]] = None, + preload_module_classes: Optional[list[str]] = None, + tied_params_map: Optional[dict[int, dict[torch.device, torch.Tensor]]] = None, +): + """ + Recursively attaches `AlignDevicesHook` to all submodules of a given model to make sure they have the right + execution device + + Args: + module (`torch.nn.Module`): + The module where we want to attach the hooks. + execution_device (`int`, `str` or `torch.device`): + The device on which inputs and model weights should be placed before the forward pass. + skip_keys (`str` or `List[str]`, *optional*): + A list of keys to ignore when moving inputs or outputs between devices. + preload_module_classes (`List[str]`, *optional*): + A list of classes whose instances should load all their weights (even in the submodules) at the beginning + of the forward. This should only be used for classes that have submodules which are registered but not + called directly during the forward, for instance if a `dense` linear layer is registered, but at forward, + `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly. + tied_params_map (Optional[Dict[int, Dict[torch.device, torch.Tensor]]], *optional*, defaults to `None`): + A map of data pointers to dictionaries of devices to already dispatched tied weights. For a given execution + device, this parameter is useful to reuse the first available pointer of a shared weight for all others, + instead of duplicating memory. + """ + if not hasattr(module, "_hf_hook") and len(module.state_dict()) > 0: + add_hook_to_module( + module, + AlignDevicesHook(execution_device, skip_keys=skip_keys, tied_params_map=tied_params_map), + ) + + # Break the recursion if we get to a preload module. + if preload_module_classes is not None and module.__class__.__name__ in preload_module_classes: + return + + for child in module.children(): + attach_execution_device_hook( + child, + execution_device, + skip_keys=skip_keys, + preload_module_classes=preload_module_classes, + tied_params_map=tied_params_map, + ) + + +def attach_align_device_hook( + module: torch.nn.Module, + execution_device: Optional[torch.device] = None, + offload: bool = False, + weights_map: Optional[Mapping] = None, + offload_buffers: bool = False, + module_name: str = "", + skip_keys: Optional[Union[str, list[str]]] = None, + preload_module_classes: Optional[list[str]] = None, + tied_params_map: Optional[dict[int, dict[torch.device, torch.Tensor]]] = None, +): + """ + Recursively attaches `AlignDevicesHook` to all submodules of a given model that have direct parameters and/or + buffers. + + Args: + module (`torch.nn.Module`): + The module where we want to attach the hooks. + execution_device (`torch.device`, *optional*): + The device on which inputs and model weights should be placed before the forward pass. + offload (`bool`, *optional*, defaults to `False`): + Whether or not the weights should be offloaded after the forward pass. + weights_map (`Mapping[str, torch.Tensor]`, *optional*): + When the model weights are offloaded, a (potentially lazy) map from param names to the tensor values. + offload_buffers (`bool`, *optional*, defaults to `False`): + Whether or not to include the associated module's buffers when offloading. + module_name (`str`, *optional*, defaults to `""`): + The name of the module. + skip_keys (`str` or `List[str]`, *optional*): + A list of keys to ignore when moving inputs or outputs between devices. + preload_module_classes (`List[str]`, *optional*): + A list of classes whose instances should load all their weights (even in the submodules) at the beginning + of the forward. This should only be used for classes that have submodules which are registered but not + called directly during the forward, for instance if a `dense` linear layer is registered, but at forward, + `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly. + tied_params_map (Optional[Dict[int, Dict[torch.device, torch.Tensor]]], *optional*, defaults to `None`): + A map of data pointers to dictionaries of devices to already dispatched tied weights. For a given execution + device, this parameter is useful to reuse the first available pointer of a shared weight for all others, + instead of duplicating memory. + """ + # Attach the hook on this module if it has any direct tensor. + directs = named_module_tensors(module) + full_offload = ( + offload and preload_module_classes is not None and module.__class__.__name__ in preload_module_classes + ) + + if len(list(directs)) > 0 or full_offload: + if weights_map is not None: + prefix = f"{module_name}." if len(module_name) > 0 else "" + prefixed_weights_map = PrefixedDataset(weights_map, prefix) + else: + prefixed_weights_map = None + hook = AlignDevicesHook( + execution_device=execution_device, + offload=offload, + weights_map=prefixed_weights_map, + offload_buffers=offload_buffers, + place_submodules=full_offload, + skip_keys=skip_keys, + tied_params_map=tied_params_map, + ) + add_hook_to_module(module, hook, append=True) + + # We stop the recursion in case we hit the full offload. + if full_offload: + return + + # Recurse on all children of the module. + for child_name, child in module.named_children(): + child_name = f"{module_name}.{child_name}" if len(module_name) > 0 else child_name + attach_align_device_hook( + child, + execution_device=execution_device, + offload=offload, + weights_map=weights_map, + offload_buffers=offload_buffers, + module_name=child_name, + preload_module_classes=preload_module_classes, + skip_keys=skip_keys, + tied_params_map=tied_params_map, + ) + + +def remove_hook_from_submodules(module: nn.Module): + """ + Recursively removes all hooks attached on the submodules of a given model. + + Args: + module (`torch.nn.Module`): The module on which to remove all hooks. + """ + remove_hook_from_module(module) + for child in module.children(): + remove_hook_from_submodules(child) + + +def attach_align_device_hook_on_blocks( + module: nn.Module, + execution_device: Optional[Union[torch.device, dict[str, torch.device]]] = None, + offload: Union[bool, dict[str, bool]] = False, + weights_map: Optional[Mapping] = None, + offload_buffers: bool = False, + module_name: str = "", + skip_keys: Optional[Union[str, list[str]]] = None, + preload_module_classes: Optional[list[str]] = None, + tied_params_map: Optional[dict[int, dict[torch.device, torch.Tensor]]] = None, +): + """ + Attaches `AlignDevicesHook` to all blocks of a given model as needed. + + Args: + module (`torch.nn.Module`): + The module where we want to attach the hooks. + execution_device (`torch.device` or `Dict[str, torch.device]`, *optional*): + The device on which inputs and model weights should be placed before the forward pass. It can be one device + for the whole module, or a dictionary mapping module name to device. + offload (`bool`, *optional*, defaults to `False`): + Whether or not the weights should be offloaded after the forward pass. It can be one boolean for the whole + module, or a dictionary mapping module name to boolean. + weights_map (`Mapping[str, torch.Tensor]`, *optional*): + When the model weights are offloaded, a (potentially lazy) map from param names to the tensor values. + offload_buffers (`bool`, *optional*, defaults to `False`): + Whether or not to include the associated module's buffers when offloading. + module_name (`str`, *optional*, defaults to `""`): + The name of the module. + skip_keys (`str` or `List[str]`, *optional*): + A list of keys to ignore when moving inputs or outputs between devices. + preload_module_classes (`List[str]`, *optional*): + A list of classes whose instances should load all their weights (even in the submodules) at the beginning + of the forward. This should only be used for classes that have submodules which are registered but not + called directly during the forward, for instance if a `dense` linear layer is registered, but at forward, + `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly. + tied_params_map (Optional[Dict[int, Dict[torch.device, torch.Tensor]]], *optional*, defaults to `None`): + A map of data pointers to dictionaries of devices to already dispatched tied weights. For a given execution + device, this parameter is useful to reuse the first available pointer of a shared weight for all others, + instead of duplicating memory. + """ + # If one device and one offload, we've got one hook. + if not isinstance(execution_device, Mapping) and not isinstance(offload, dict): + if not offload: + hook = AlignDevicesHook( + execution_device=execution_device, + io_same_device=True, + skip_keys=skip_keys, + place_submodules=True, + tied_params_map=tied_params_map, + ) + add_hook_to_module(module, hook) + else: + attach_align_device_hook( + module, + execution_device=execution_device, + offload=True, + weights_map=weights_map, + offload_buffers=offload_buffers, + module_name=module_name, + skip_keys=skip_keys, + tied_params_map=tied_params_map, + ) + return + + if not isinstance(execution_device, Mapping): + execution_device = {key: execution_device for key in offload.keys()} + if not isinstance(offload, Mapping): + offload = {key: offload for key in execution_device.keys()} + + if module_name in execution_device and module_name in offload and not offload[module_name]: + hook = AlignDevicesHook( + execution_device=execution_device[module_name], + offload_buffers=offload_buffers, + io_same_device=(module_name == ""), + place_submodules=True, + skip_keys=skip_keys, + tied_params_map=tied_params_map, + ) + add_hook_to_module(module, hook) + attach_execution_device_hook( + module, execution_device[module_name], skip_keys=skip_keys, tied_params_map=tied_params_map + ) + elif module_name in execution_device and module_name in offload: + attach_align_device_hook( + module, + execution_device=execution_device[module_name], + offload=True, + weights_map=weights_map, + offload_buffers=offload_buffers, + module_name=module_name, + skip_keys=skip_keys, + preload_module_classes=preload_module_classes, + tied_params_map=tied_params_map, + ) + if not hasattr(module, "_hf_hook"): + hook = AlignDevicesHook( + execution_device=execution_device[module_name], + io_same_device=(module_name == ""), + skip_keys=skip_keys, + tied_params_map=tied_params_map, + ) + add_hook_to_module(module, hook) + attach_execution_device_hook( + module, + execution_device[module_name], + preload_module_classes=preload_module_classes, + skip_keys=skip_keys, + tied_params_map=tied_params_map, + ) + elif module_name == "": + hook = AlignDevicesHook( + execution_device=execution_device.get(""), + io_same_device=True, + skip_keys=skip_keys, + tied_params_map=tied_params_map, + ) + add_hook_to_module(module, hook) + + for child_name, child in module.named_children(): + child_name = f"{module_name}.{child_name}" if len(module_name) > 0 else child_name + attach_align_device_hook_on_blocks( + child, + execution_device=execution_device, + offload=offload, + weights_map=weights_map, + offload_buffers=offload_buffers, + module_name=child_name, + preload_module_classes=preload_module_classes, + skip_keys=skip_keys, + tied_params_map=tied_params_map, + ) + + +class CpuOffload(ModelHook): + """ + Offloads a model on the CPU until its forward pass is called. The model will not be offloaded back to the CPU after + the forward, the user needs to call the `init_hook` method again for this. + + Args: + execution_device(`str`, `int` or `torch.device`, *optional*): + The device on which the model should be executed. Will default to the MPS device if it's available, then + GPU 0 if there is a GPU, and finally to the CPU. + prev_module_hook (`UserCpuOffloadHook`, *optional*): + The hook sent back by [`cpu_offload_with_hook`] for a previous model in the pipeline you are running. If + passed, its offload method will be called just before the forward of the model to which this hook is + attached. + """ + + def __init__( + self, + execution_device: Optional[Union[str, int, torch.device]] = None, + prev_module_hook: Optional["UserCpuOffloadHook"] = None, + ): + self.prev_module_hook = prev_module_hook + + self.execution_device = execution_device if execution_device is not None else PartialState().default_device + + def init_hook(self, module): + return module.to("cpu") + + def pre_forward(self, module, *args, **kwargs): + if self.prev_module_hook is not None and isinstance(self.prev_module_hook, UserCpuOffloadHook): + prev_module = self.prev_module_hook.model + prev_device = next(prev_module.parameters()).device + + # Only offload the previous module if it is not already on CPU. + if prev_device != torch.device("cpu"): + self.prev_module_hook.offload() + clear_device_cache() + + # If the current device is already the self.execution_device, we can skip the transfer. + current_device = next(module.parameters()).device + if current_device == self.execution_device: + return args, kwargs + + module.to(self.execution_device) + return send_to_device(args, self.execution_device), send_to_device(kwargs, self.execution_device) + + +class UserCpuOffloadHook: + """ + A simple hook grouping a model and a `ModelHook`, which provides easy APIs for to call the init method of the hook + or remove it entirely. + """ + + def __init__(self, model, hook): + self.model = model + self.hook = hook + + def offload(self): + self.hook.init_hook(self.model) + + def remove(self): + remove_hook_from_module(self.model) + + +class LayerwiseCastingHook(ModelHook): + r""" + A hook that casts the weights of a module to a high precision dtype for computation, and to a low precision dtype + for storage. This process may lead to quality loss in the output, but can significantly reduce the memory + footprint. + """ + + _is_stateful = False + + def __init__(self, storage_dtype: torch.dtype, compute_dtype: torch.dtype, non_blocking: bool) -> None: + self.storage_dtype = storage_dtype + self.compute_dtype = compute_dtype + self.non_blocking = non_blocking + + def init_hook(self, module: torch.nn.Module): + module.to(dtype=self.storage_dtype, non_blocking=self.non_blocking) + return module + + def pre_forward(self, module: torch.nn.Module, *args, **kwargs): + module.to(dtype=self.compute_dtype, non_blocking=self.non_blocking) + return args, kwargs + + def post_forward(self, module: torch.nn.Module, output): + module.to(dtype=self.storage_dtype, non_blocking=self.non_blocking) + return output diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/inference.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/inference.py new file mode 100644 index 0000000000000000000000000000000000000000..35c99399baba62f77bbaac8adf4e2eacdbabf779 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/inference.py @@ -0,0 +1,186 @@ +# Copyright 2024 The HuggingFace Team. All rights reserved. +# +# 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. +import math +from types import MethodType +from typing import Any, Optional, Union + +from .state import PartialState +from .utils import ( + calculate_maximum_sizes, + convert_bytes, + copy_tensor_to_devices, + ignorant_find_batch_size, + infer_auto_device_map, + is_pippy_available, + pad_input_tensors, + send_to_device, +) + + +def generate_device_map( + model, num_processes: int = 1, no_split_module_classes=None, max_memory: Optional[dict] = None +): + """ + Calculates the device map for `model` with an offset for PiPPy + """ + if num_processes == 1: + return infer_auto_device_map(model, no_split_module_classes=no_split_module_classes, clean_result=False) + if max_memory is None: + model_size, shared = calculate_maximum_sizes(model) + + # Split into `n` chunks for each GPU + memory = (model_size + shared[0]) / num_processes + memory = convert_bytes(memory) + value, ending = memory.split(" ") + + # Add a chunk to deal with potential extra shared memory instances + memory = math.ceil(float(value)) * 1.1 + memory = f"{memory} {ending}" + max_memory = {i: memory for i in range(num_processes)} + device_map = infer_auto_device_map( + model, + max_memory=max_memory, + no_split_module_classes=no_split_module_classes, + clean_result=False, + ) + return device_map + + +def find_pippy_batch_size(args, kwargs): + found_batch_size = None + if args is not None: + for arg in args: + found_batch_size = ignorant_find_batch_size(arg) + if found_batch_size is not None: + break + if kwargs is not None and found_batch_size is None: + for kwarg in kwargs.values(): + found_batch_size = ignorant_find_batch_size(kwarg) + if found_batch_size is not None: + break + return found_batch_size + + +def build_pipeline(model, split_points, args, kwargs, num_chunks): + """ + Attaches the split points to the model based on `self.device_map` and generates a `PipelineStage`. Requires passing + in needed `args` and `kwargs` as the model needs on the CPU. + + Users can pass in custom `num_chunks` as an optional hyper-parameter. By default will use + `AcceleratorState.num_processes` + """ + # Note: We import here to reduce import time from general modules, and isolate outside dependencies + from torch.distributed.pipelining import ScheduleGPipe, SplitPoint, pipeline + + # We need to annotate the split points in the model for PiPPy + state = PartialState() + split_spec = {split_point: SplitPoint.BEGINNING for split_point in split_points} + pipe = pipeline( + model, + mb_args=args, + mb_kwargs=kwargs, + split_spec=split_spec, + ) + stage = pipe.build_stage(state.local_process_index, device=state.device) + schedule = ScheduleGPipe(stage, num_chunks) + + return schedule + + +def pippy_forward(forward, num_chunks, gather_output, *args, **kwargs): + state = PartialState() + output = None + + if state.num_processes == 1: + output = forward(*args, **kwargs) + elif state.is_local_main_process: + found_batch_size = find_pippy_batch_size(args, kwargs) + if found_batch_size is None: + raise ValueError("Could not find batch size from args or kwargs") + else: + if found_batch_size != num_chunks: + args = pad_input_tensors(args, found_batch_size, num_chunks) + kwargs = pad_input_tensors(kwargs, found_batch_size, num_chunks) + forward(*args, **kwargs) + elif state.is_last_process: + output = forward() + else: + forward() + if gather_output: + # Each node will get a copy of the full output which is only on the last GPU + output = copy_tensor_to_devices(output) + return output + + +def prepare_pippy( + model, + split_points: Optional[Union[str, list[str]]] = "auto", + no_split_module_classes: Optional[list[str]] = None, + example_args: Optional[tuple[Any]] = (), + example_kwargs: Optional[dict[str, Any]] = None, + num_chunks: Optional[int] = None, + gather_output: Optional[bool] = False, +): + """ + Wraps `model` for pipeline parallel inference. + + Args: + model (`torch.nn.Module`): + A model we want to split for pipeline-parallel inference + split_points (`str` or `List[str]`, defaults to 'auto'): + How to generate the split points and chunk the model across each GPU. 'auto' will find the best balanced + split given any model. Should be a list of layer names in the model to split by otherwise. + no_split_module_classes (`List[str]`): + A list of class names for layers we don't want to be split. + example_args (tuple of model inputs): + The expected inputs for the model that uses order-based inputs for a *single process*. Recommended to use + this method if possible. + example_kwargs (dict of model inputs) + The expected inputs for the model that uses dictionary-based inputs for a *single process*. This is a + *highly* limiting structure that requires the same keys be present at *all* inference calls. Not + recommended unless the prior condition is true for all cases. + num_chunks (`int`, defaults to the number of available GPUs): + The number of different stages the Pipeline will have. By default it will assign one chunk per GPU, but + this can be tuned and played with. In general one should have num_chunks >= num_gpus. + gather_output (`bool`, defaults to `False`): + If `True`, the output from the last GPU (which holds the true outputs) is sent across to all GPUs. + """ + if not is_pippy_available(): + raise ImportError("Using `torch.distributed.pipelining` requires PyTorch 2.4.0 or later.") + state = PartialState() + example_args = send_to_device(example_args, "cpu") + example_kwargs = send_to_device(example_kwargs, "cpu") + if num_chunks is None: + num_chunks = state.num_processes + if split_points == "auto": + device_map = generate_device_map(model, num_chunks, no_split_module_classes=no_split_module_classes) + split_points = [] + for i in range(1, num_chunks): + split_points.append(next(k for k, v in device_map.items() if v == i)) + model.hf_split_points = split_points + stage = build_pipeline(model, split_points, example_args, example_kwargs, num_chunks) + model._original_forward = model.forward + model._original_call = model.__call__ + model.pippy_stage = stage + model.hf_split_points = split_points + + def forward(*args, **kwargs): + return pippy_forward(stage.step, num_chunks, gather_output, *args, **kwargs) + + # To act like a decorator so that it can be popped when doing `extract_model_from_parallel` + # Note: creates an infinite recursion loop with `generate` + model_forward = MethodType(forward, model) + forward.__wrapped__ = model_forward + model.forward = forward + return model diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/launchers.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/launchers.py new file mode 100644 index 0000000000000000000000000000000000000000..d811c1622e8f3fc4d7535902af54274a49b5b03d --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/launchers.py @@ -0,0 +1,307 @@ +# Copyright 2021 The HuggingFace Team. All rights reserved. +# +# 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. + +import os +import sys +import tempfile + +import torch + +from .state import AcceleratorState, PartialState +from .utils import ( + PrecisionType, + PrepareForLaunch, + are_libraries_initialized, + check_cuda_p2p_ib_support, + get_current_device_type, + get_gpu_info, + is_mps_available, + is_torch_version, + patch_environment, +) +from .utils.constants import ELASTIC_LOG_LINE_PREFIX_TEMPLATE_PYTORCH_VERSION + + +def test_launch(): + "Verify a `PartialState` can be initialized." + _ = PartialState() + + +def notebook_launcher( + function, + args=(), + num_processes=None, + mixed_precision="no", + use_port="29500", + master_addr="127.0.0.1", + node_rank=0, + num_nodes=1, + rdzv_backend="static", + rdzv_endpoint="", + rdzv_conf=None, + rdzv_id="none", + max_restarts=0, + monitor_interval=0.1, + log_line_prefix_template=None, +): + """ + Launches a training function, using several processes or multiple nodes if it's possible in the current environment + (TPU with multiple cores for instance). + + + + To use this function absolutely zero calls to a device must be made in the notebook session before calling. If any + have been made, you will need to restart the notebook and make sure no cells use any device capability. + + Setting `ACCELERATE_DEBUG_MODE="1"` in your environment will run a test before truly launching to ensure that none + of those calls have been made. + + + + Args: + function (`Callable`): + The training function to execute. If it accepts arguments, the first argument should be the index of the + process run. + args (`Tuple`): + Tuple of arguments to pass to the function (it will receive `*args`). + num_processes (`int`, *optional*): + The number of processes to use for training. Will default to 8 in Colab/Kaggle if a TPU is available, to + the number of devices available otherwise. + mixed_precision (`str`, *optional*, defaults to `"no"`): + If `fp16` or `bf16`, will use mixed precision training on multi-device. + use_port (`str`, *optional*, defaults to `"29500"`): + The port to use to communicate between processes when launching a multi-device training. + master_addr (`str`, *optional*, defaults to `"127.0.0.1"`): + The address to use for communication between processes. + node_rank (`int`, *optional*, defaults to 0): + The rank of the current node. + num_nodes (`int`, *optional*, defaults to 1): + The number of nodes to use for training. + rdzv_backend (`str`, *optional*, defaults to `"static"`): + The rendezvous method to use, such as 'static' (the default) or 'c10d' + rdzv_endpoint (`str`, *optional*, defaults to `""`): + The endpoint of the rdzv sync. storage. + rdzv_conf (`Dict`, *optional*, defaults to `None`): + Additional rendezvous configuration. + rdzv_id (`str`, *optional*, defaults to `"none"`): + The unique run id of the job. + max_restarts (`int`, *optional*, defaults to 0): + The maximum amount of restarts that elastic agent will conduct on workers before failure. + monitor_interval (`float`, *optional*, defaults to 0.1): + The interval in seconds that is used by the elastic_agent as a period of monitoring workers. + log_line_prefix_template (`str`, *optional*, defaults to `None`): + The prefix template for elastic launch logging. Available from PyTorch 2.2.0. + + Example: + + ```python + # Assume this is defined in a Jupyter Notebook on an instance with two devices + from accelerate import notebook_launcher + + + def train(*args): + # Your training function here + ... + + + notebook_launcher(train, args=(arg1, arg2), num_processes=2, mixed_precision="fp16") + ``` + """ + # Are we in a google colab or a Kaggle Kernel? + in_colab = False + in_kaggle = False + if any(key.startswith("KAGGLE") for key in os.environ.keys()): + in_kaggle = True + elif "IPython" in sys.modules: + in_colab = "google.colab" in str(sys.modules["IPython"].get_ipython()) + + try: + mixed_precision = PrecisionType(mixed_precision.lower()) + except ValueError: + raise ValueError( + f"Unknown mixed_precision mode: {args.mixed_precision.lower()}. Choose between {PrecisionType.list()}." + ) + + if (in_colab or in_kaggle) and ( + (os.environ.get("TPU_NAME", None) is not None) or (os.environ.get("PJRT_DEVICE", "") == "TPU") + ): + # TPU launch + import torch_xla.distributed.xla_multiprocessing as xmp + + if len(AcceleratorState._shared_state) > 0: + raise ValueError( + "To train on TPU in Colab or Kaggle Kernel, the `Accelerator` should only be initialized inside " + "your training function. Restart your notebook and make sure no cells initializes an " + "`Accelerator`." + ) + + launcher = PrepareForLaunch(function, distributed_type="XLA") + print("Launching a training on TPU cores.") + xmp.spawn(launcher, args=args, start_method="fork") + elif in_colab and get_gpu_info()[1] < 2: + # No need for a distributed launch otherwise as it's either CPU or one GPU. + if torch.cuda.is_available(): + print("Launching training on one GPU.") + else: + print("Launching training on one CPU.") + function(*args) + else: + if num_processes is None: + raise ValueError( + "You have to specify the number of devices you would like to use, add `num_processes=...` to your call." + ) + if node_rank >= num_nodes: + raise ValueError("The node_rank must be less than the number of nodes.") + if num_processes > 1: + # Multi-device launch + from torch.distributed.launcher.api import LaunchConfig, elastic_launch + from torch.multiprocessing import start_processes + from torch.multiprocessing.spawn import ProcessRaisedException + + if len(AcceleratorState._shared_state) > 0: + raise ValueError( + "To launch a multi-device training from your notebook, the `Accelerator` should only be initialized " + "inside your training function. Restart your notebook and make sure no cells initializes an " + "`Accelerator`." + ) + # Check for specific libraries known to initialize device that users constantly use + problematic_imports = are_libraries_initialized("bitsandbytes") + if len(problematic_imports) > 0: + err = ( + "Could not start distributed process. Libraries known to initialize device upon import have been " + "imported already. Please keep these imports inside your training function to try and help with this:" + ) + for lib_name in problematic_imports: + err += f"\n\t* `{lib_name}`" + raise RuntimeError(err) + + patched_env = dict( + nproc=num_processes, + node_rank=node_rank, + world_size=num_nodes * num_processes, + master_addr=master_addr, + master_port=use_port, + mixed_precision=mixed_precision, + ) + + # Check for CUDA P2P and IB issues + if not check_cuda_p2p_ib_support(): + patched_env["nccl_p2p_disable"] = "1" + patched_env["nccl_ib_disable"] = "1" + + # torch.distributed will expect a few environment variable to be here. We set the ones common to each + # process here (the other ones will be set be the launcher). + with patch_environment(**patched_env): + # First dummy launch + # Determine device type without initializing any device (which would break fork) + device_type, distributed_type = get_current_device_type() + if os.environ.get("ACCELERATE_DEBUG_MODE", "false").lower() == "true": + launcher = PrepareForLaunch(test_launch, distributed_type=distributed_type) + try: + start_processes(launcher, args=(), nprocs=num_processes, start_method="fork") + except ProcessRaisedException as e: + err = "An issue was found when verifying a stable environment for the notebook launcher." + if f"Cannot re-initialize {device_type.upper()} in forked subprocess" in e.args[0]: + raise RuntimeError( + f"{err}" + "This likely stems from an outside import causing issues once the `notebook_launcher()` is called. " + "Please review your imports and test them when running the `notebook_launcher()` to identify " + f"which one is problematic and causing {device_type.upper()} to be initialized." + ) from e + else: + raise RuntimeError(f"{err} The following error was raised: {e}") from e + # Now the actual launch + launcher = PrepareForLaunch(function, distributed_type=distributed_type) + print(f"Launching training on {num_processes} {device_type.upper()}s.") + try: + if rdzv_conf is None: + rdzv_conf = {} + if rdzv_backend == "static": + rdzv_conf["rank"] = node_rank + if not rdzv_endpoint: + rdzv_endpoint = f"{master_addr}:{use_port}" + launch_config_kwargs = dict( + min_nodes=num_nodes, + max_nodes=num_nodes, + nproc_per_node=num_processes, + run_id=rdzv_id, + rdzv_endpoint=rdzv_endpoint, + rdzv_backend=rdzv_backend, + rdzv_configs=rdzv_conf, + max_restarts=max_restarts, + monitor_interval=monitor_interval, + start_method="fork", + ) + if is_torch_version(">=", ELASTIC_LOG_LINE_PREFIX_TEMPLATE_PYTORCH_VERSION): + launch_config_kwargs["log_line_prefix_template"] = log_line_prefix_template + elastic_launch(config=LaunchConfig(**launch_config_kwargs), entrypoint=function)(*args) + except ProcessRaisedException as e: + if f"Cannot re-initialize {device_type.upper()} in forked subprocess" in e.args[0]: + raise RuntimeError( + f"{device_type.upper()} has been initialized before the `notebook_launcher` could create a forked subprocess. " + "This likely stems from an outside import causing issues once the `notebook_launcher()` is called. " + "Please review your imports and test them when running the `notebook_launcher()` to identify " + f"which one is problematic and causing {device_type.upper()} to be initialized." + ) from e + else: + raise RuntimeError(f"An issue was found when launching the training: {e}") from e + + else: + # No need for a distributed launch otherwise as it's either CPU, GPU, XPU or MPS. + if is_mps_available(): + os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" + print("Launching training on MPS.") + elif torch.cuda.is_available(): + print("Launching training on one GPU.") + elif torch.xpu.is_available(): + print("Launching training on one XPU.") + else: + print("Launching training on CPU.") + function(*args) + + +def debug_launcher(function, args=(), num_processes=2): + """ + Launches a training function using several processes on CPU for debugging purposes. + + + + This function is provided for internal testing and debugging, but it's not intended for real trainings. It will + only use the CPU. + + + + Args: + function (`Callable`): + The training function to execute. + args (`Tuple`): + Tuple of arguments to pass to the function (it will receive `*args`). + num_processes (`int`, *optional*, defaults to 2): + The number of processes to use for training. + """ + from torch.multiprocessing import start_processes + + with tempfile.NamedTemporaryFile() as tmp_file: + # torch.distributed will expect a few environment variable to be here. We set the ones common to each + # process here (the other ones will be set be the launcher). + with patch_environment( + world_size=num_processes, + master_addr="127.0.0.1", + master_port="29500", + accelerate_mixed_precision="no", + accelerate_debug_rdv_file=tmp_file.name, + accelerate_use_cpu="yes", + ): + launcher = PrepareForLaunch(function, debug=True) + start_processes(launcher, args=args, nprocs=num_processes, start_method="fork") diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/local_sgd.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/local_sgd.py new file mode 100644 index 0000000000000000000000000000000000000000..5913d2d03df3c67ccf3e9ef84e75e06dc694b1fa --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/local_sgd.py @@ -0,0 +1,106 @@ +# Copyright 2023 The HuggingFace Team. All rights reserved. +# +# 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. +import torch + +from accelerate import Accelerator, DistributedType + + +class LocalSGD: + """ + A helper class to support local SGD on top of Accelerator. It simply runs a given number of updates independently + on each device, and averages model weights every K synchronization step. + + It should be used only in the multi-GPU (or multi-CPU) setup without extensions such as DeepSpeed. In particular, + this is a simple implementation that cannot support scenarios such as model parallelism. + + + Although we are not aware of the true origins of this simple approach, the idea of local SGD is quite old and goes + back to at least: + + Zhang, J., De Sa, C., Mitliagkas, I., & Ré, C. (2016). [Parallel SGD: When does averaging help?. arXiv preprint + arXiv:1606.07365.](https://huggingface.co/papers/1606.07365) + + We credit the term Local SGD to the following paper (but there might be earlier references we are not aware of). + + Stich, Sebastian Urban. ["Local SGD Converges Fast and Communicates Little." ICLR 2019-International Conference on + Learning Representations. No. CONF. 2019.](https://huggingface.co/papers/1805.09767) + + """ + + def __enter__(self): + if self.enabled: + self.model_sync_obj = self.model.no_sync() + self.model_sync_obj.__enter__() + + return self + + def __exit__(self, type, value, tb): + if self.enabled: + # Average all models on exit + self._sync_and_avg_model_params() + self.model_sync_obj.__exit__(type, value, tb) + + def __init__(self, accelerator: Accelerator, model: torch.nn.Module, local_sgd_steps: int, enabled: bool = True): + """ + Constructor. + + Args: + model (`torch.nn.Module): + The model whose parameters we need to average. + accelerator (`Accelerator`): + Accelerator object. + local_sgd_steps (`int`): + A number of local SGD steps (before model parameters are synchronized). + enabled (`bool): + Local SGD is disabled if this parameter set to `False`. + """ + if accelerator.distributed_type not in [ + DistributedType.NO, + DistributedType.MULTI_CPU, + DistributedType.MULTI_GPU, + DistributedType.MULTI_XPU, + DistributedType.MULTI_MLU, + DistributedType.MULTI_HPU, + DistributedType.MULTI_SDAA, + DistributedType.MULTI_MUSA, + DistributedType.MULTI_NPU, + ]: + raise NotImplementedError("LocalSGD is supported only for CPUs and GPUs (no DeepSpeed or MegatronLM)") + self.enabled = enabled and accelerator.distributed_type != DistributedType.NO + self.num_steps = 0 + if self.enabled: + self.accelerator = accelerator + self.model = model + self.local_sgd_steps = local_sgd_steps + + def step(self): + """ + This function makes a "step" and synchronizes model parameters if necessary. + """ + self.num_steps += 1 + if not self.enabled: + return + + if self.num_steps % self.local_sgd_steps == 0: + self._sync_and_avg_model_params() + + def _sync_and_avg_model_params(self): + """ + Synchronize + Average model parameters across all GPUs + """ + + self.accelerator.wait_for_everyone() + with self.accelerator.autocast(): + for param in self.model.parameters(): + param.data = self.accelerator.reduce(param.data, reduction="mean") diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/logging.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/logging.py new file mode 100644 index 0000000000000000000000000000000000000000..172445c305b2d469c26997df35a97c1c22fe25ed --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/logging.py @@ -0,0 +1,126 @@ +# Copyright 2022 The HuggingFace Team. All rights reserved. +# +# 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. + +import functools +import logging +import os +from typing import Optional + +from .state import PartialState + + +class MultiProcessAdapter(logging.LoggerAdapter): + """ + An adapter to assist with logging in multiprocess. + + `log` takes in an additional `main_process_only` kwarg, which dictates whether it should be called on all processes + or only the main executed one. Default is `main_process_only=True`. + + Does not require an `Accelerator` object to be created first. + """ + + @staticmethod + def _should_log(main_process_only): + "Check if log should be performed" + state = PartialState() + return not main_process_only or (main_process_only and state.is_main_process) + + def log(self, level, msg, *args, **kwargs): + """ + Delegates logger call after checking if we should log. + + Accepts a new kwarg of `main_process_only`, which will dictate whether it will be logged across all processes + or only the main executed one. Default is `True` if not passed + + Also accepts "in_order", which if `True` makes the processes log one by one, in order. This is much easier to + read, but comes at the cost of sometimes needing to wait for the other processes. Default is `False` to not + break with the previous behavior. + + `in_order` is ignored if `main_process_only` is passed. + """ + if PartialState._shared_state == {}: + raise RuntimeError( + "You must initialize the accelerate state by calling either `PartialState()` or `Accelerator()` before using the logging utility." + ) + main_process_only = kwargs.pop("main_process_only", True) + in_order = kwargs.pop("in_order", False) + # set `stacklevel` to exclude ourself in `Logger.findCaller()` while respecting user's choice + kwargs.setdefault("stacklevel", 2) + + if self.isEnabledFor(level): + if self._should_log(main_process_only): + msg, kwargs = self.process(msg, kwargs) + self.logger.log(level, msg, *args, **kwargs) + + elif in_order: + state = PartialState() + for i in range(state.num_processes): + if i == state.process_index: + msg, kwargs = self.process(msg, kwargs) + self.logger.log(level, msg, *args, **kwargs) + state.wait_for_everyone() + + @functools.lru_cache(None) + def warning_once(self, *args, **kwargs): + """ + This method is identical to `logger.warning()`, but will emit the warning with the same message only once + + Note: The cache is for the function arguments, so 2 different callers using the same arguments will hit the + cache. The assumption here is that all warning messages are unique across the code. If they aren't then need to + switch to another type of cache that includes the caller frame information in the hashing function. + """ + self.warning(*args, **kwargs) + + +def get_logger(name: str, log_level: Optional[str] = None): + """ + Returns a `logging.Logger` for `name` that can handle multiprocessing. + + If a log should be called on all processes, pass `main_process_only=False` If a log should be called on all + processes and in order, also pass `in_order=True` + + Args: + name (`str`): + The name for the logger, such as `__file__` + log_level (`str`, *optional*): + The log level to use. If not passed, will default to the `LOG_LEVEL` environment variable, or `INFO` if not + + Example: + + ```python + >>> from accelerate.logging import get_logger + >>> from accelerate import Accelerator + + >>> logger = get_logger(__name__) + + >>> accelerator = Accelerator() + >>> logger.info("My log", main_process_only=False) + >>> logger.debug("My log", main_process_only=True) + + >>> logger = get_logger(__name__, log_level="DEBUG") + >>> logger.info("My log") + >>> logger.debug("My second log") + + >>> array = ["a", "b", "c", "d"] + >>> letter_at_rank = array[accelerator.process_index] + >>> logger.info(letter_at_rank, in_order=True) + ``` + """ + if log_level is None: + log_level = os.environ.get("ACCELERATE_LOG_LEVEL", None) + logger = logging.getLogger(name) + if log_level is not None: + logger.setLevel(log_level.upper()) + logger.root.setLevel(log_level.upper()) + return MultiProcessAdapter(logger, {}) diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/memory_utils.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/memory_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..fa2e2c8b9d7d0064c3e5e282737a7ad6919bde29 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/memory_utils.py @@ -0,0 +1,22 @@ +# Copyright 2022 The HuggingFace Team. All rights reserved. +# +# 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. + +import warnings + + +warnings.warn( + "memory_utils has been reorganized to utils.memory. Import `find_executable_batchsize` from the main `__init__`: " + "`from accelerate import find_executable_batch_size` to avoid this warning.", + FutureWarning, +) diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/optimizer.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/optimizer.py new file mode 100644 index 0000000000000000000000000000000000000000..c1f8faa15436fb5d1d02ae4d691a4fcac44cd236 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/optimizer.py @@ -0,0 +1,213 @@ +# Copyright 2021 The HuggingFace Team. All rights reserved. +# +# 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. + +import inspect + +import torch + +from .state import AcceleratorState, GradientState +from .utils import DistributedType, honor_type, is_lomo_available, is_torch_xla_available + + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + import torch_xla.runtime as xr + + +def move_to_device(state, device): + if isinstance(state, (list, tuple)): + return honor_type(state, (move_to_device(t, device) for t in state)) + elif isinstance(state, dict): + return type(state)({k: move_to_device(v, device) for k, v in state.items()}) + elif isinstance(state, torch.Tensor): + return state.to(device) + return state + + +class AcceleratedOptimizer(torch.optim.Optimizer): + """ + Internal wrapper around a torch optimizer. + + Conditionally will perform `step` and `zero_grad` if gradients should be synchronized when performing gradient + accumulation. + + Args: + optimizer (`torch.optim.optimizer.Optimizer`): + The optimizer to wrap. + device_placement (`bool`, *optional*, defaults to `True`): + Whether or not the optimizer should handle device placement. If so, it will place the state dictionary of + `optimizer` on the right device. + scaler (`torch.amp.GradScaler` or `torch.cuda.amp.GradScaler`, *optional*): + The scaler to use in the step function if training with mixed precision. + """ + + def __init__(self, optimizer, device_placement=True, scaler=None): + self.optimizer = optimizer + self.scaler = scaler + self.accelerator_state = AcceleratorState() + self.gradient_state = GradientState() + self.device_placement = device_placement + self._is_overflow = False + + if self.scaler is not None: + self._accelerate_step_called = False + self._optimizer_original_step_method = self.optimizer.step + self._optimizer_patched_step_method = patch_optimizer_step(self, self.optimizer.step) + + # Handle device placement + if device_placement: + state_dict = self.optimizer.state_dict() + if self.accelerator_state.distributed_type == DistributedType.XLA: + xm.send_cpu_data_to_device(state_dict, self.accelerator_state.device) + else: + state_dict = move_to_device(state_dict, self.accelerator_state.device) + self.optimizer.load_state_dict(state_dict) + + @property + def state(self): + return self.optimizer.state + + @state.setter + def state(self, state): + self.optimizer.state = state + + @property + def param_groups(self): + return self.optimizer.param_groups + + @param_groups.setter + def param_groups(self, param_groups): + self.optimizer.param_groups = param_groups + + @property + def defaults(self): + return self.optimizer.defaults + + @defaults.setter + def defaults(self, defaults): + self.optimizer.defaults = defaults + + def add_param_group(self, param_group): + self.optimizer.add_param_group(param_group) + + def load_state_dict(self, state_dict): + if self.accelerator_state.distributed_type == DistributedType.XLA and self.device_placement: + xm.send_cpu_data_to_device(state_dict, self.accelerator_state.device) + self.optimizer.load_state_dict(state_dict) + + def state_dict(self): + return self.optimizer.state_dict() + + def zero_grad(self, set_to_none=None): + if self.gradient_state.sync_gradients: + accept_arg = "set_to_none" in inspect.signature(self.optimizer.zero_grad).parameters + if accept_arg: + if set_to_none is None: + set_to_none = True + self.optimizer.zero_grad(set_to_none=set_to_none) + else: + if set_to_none is not None: + raise ValueError("`set_to_none` for Optimizer.zero_grad` is not supported by this optimizer.") + self.optimizer.zero_grad() + + def train(self): + """ + Sets the optimizer to "train" mode. Useful for optimizers like `schedule_free` + """ + if hasattr(self.optimizer, "train") and callable(self.optimizer.train): + self.optimizer.train() + elif ( + hasattr(self.optimizer, "optimizer") + and hasattr(self.optimizer.optimizer, "train") + and callable(self.optimizer.optimizer.train) + ): + # the deepspeed optimizer further wraps the optimizer + self.optimizer.optimizer.train() + + def eval(self): + """ + Sets the optimizer to "eval" mode. Useful for optimizers like `schedule_free` + """ + if hasattr(self.optimizer, "eval") and callable(self.optimizer.eval): + self.optimizer.eval() + + def step(self, closure=None): + if is_lomo_available(): + from lomo_optim import AdaLomo, Lomo + + if ( + not self.gradient_state.is_xla_gradients_synced + and self.accelerator_state.distributed_type == DistributedType.XLA + ): + gradients = xm._fetch_gradients(self.optimizer) + xm.all_reduce("sum", gradients, scale=1.0 / xr.world_size()) + self.gradient_state.is_xla_gradients_synced = True + + if is_lomo_available(): + # `step` should be a no-op for LOMO optimizers. + if isinstance(self.optimizer, (Lomo, AdaLomo)): + return + + if self.gradient_state.sync_gradients: + if self.scaler is not None: + self.optimizer.step = self._optimizer_patched_step_method + + self.scaler.step(self.optimizer, closure) + self.scaler.update() + + if not self._accelerate_step_called: + # If the optimizer step was skipped, gradient overflow was detected. + self._is_overflow = True + else: + self._is_overflow = False + # Reset the step method to the original one + self.optimizer.step = self._optimizer_original_step_method + # Reset the indicator + self._accelerate_step_called = False + else: + self.optimizer.step(closure) + if self.accelerator_state.distributed_type == DistributedType.XLA: + self.gradient_state.is_xla_gradients_synced = False + + def _switch_parameters(self, parameters_map): + for param_group in self.optimizer.param_groups: + param_group["params"] = [parameters_map.get(p, p) for p in param_group["params"]] + + @property + def step_was_skipped(self): + """Whether or not the optimizer step was skipped.""" + return self._is_overflow + + def __getstate__(self): + _ignored_keys = [ + "_accelerate_step_called", + "_optimizer_original_step_method", + "_optimizer_patched_step_method", + ] + return {k: v for k, v in self.__dict__.items() if k not in _ignored_keys} + + def __setstate__(self, state): + self.__dict__.update(state) + if self.scaler is not None: + self._accelerate_step_called = False + self._optimizer_original_step_method = self.optimizer.step + self._optimizer_patched_step_method = patch_optimizer_step(self, self.optimizer.step) + + +def patch_optimizer_step(accelerated_optimizer: AcceleratedOptimizer, method): + def patched_step(*args, **kwargs): + accelerated_optimizer._accelerate_step_called = True + return method(*args, **kwargs) + + return patched_step diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/parallelism_config.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/parallelism_config.py new file mode 100644 index 0000000000000000000000000000000000000000..0869d0aca232407f1891e2d8109eb33026c9c9d2 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/parallelism_config.py @@ -0,0 +1,378 @@ +# +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# 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. + +import os +import warnings +from dataclasses import dataclass +from typing import TYPE_CHECKING, Literal, Optional, Union + +from accelerate.utils.dataclasses import ( + DeepSpeedSequenceParallelConfig, + DistributedType, + TorchContextParallelConfig, + TorchTensorParallelConfig, +) +from accelerate.utils.versions import is_torch_version + + +if TYPE_CHECKING: + from accelerate import Accelerator + + +@dataclass +class ParallelismConfig: + """ + A dataclass to configure parallelisms applied to the model. Inspired by torchtitan's `ParallelDims` + https://github.com/pytorch/torchtitan/blob/main/torchtitan/distributed/parallel_dims.py + + Args: + dp_replicate_size (`int`, defaults to `1`): + The size of the data parallel group. If `dp_replicate_size` is set to 1, the data parallel replication + group will not be used. + dp_shard_size (`int`, defaults to `1`): + The size of the model shard group. If `dp_replicate_size > 1` and `tp_size > 1`, `dp_shard_size` must also + be greater than 1, as composing DDP + TP is currently not supported. + tp_size (`int`, defaults to `1`): + The size of the tensor parallel group. If `tp_size` is set to `1`, the tensor parallel group will not be + used. + tp_handler (`~utils.TorchTensorParallelConfig`, defaults to `None`): + The handler for the tensor parallel group. + cp_size (`int`, defaults to `1`): + The size of the context parallel group. Currently not supported, but reserved for future use and enabled + for downstream libraries. + cp_backend (`str`, defaults to `torch`): + Which CP backend to use: `torch` (FSDP2) + sp_size (`int`, defaults to `1`): + The size of the sequence parallel group. + sp_backend (`str`, defaults to `deepspeed`): + Which SP backend to use:`deepspeed` (ALST/Ulysses) + + You may obtain different distributed data parallel paradigms by configuring `dp_replicate_size` and `dp_shard_size` + together: + - `dp_replicate_size == 1` and `dp_shard_size > 1`, we obtain Fully Sharded Data Parallel (FSDP). + - `dp_replicate_size > 1` and `dp_shard_size > 1`, we obtain Hybrid Sharded Data Parallel (HSDP). + - `dp_replicate_size > 1` and `dp_shard_size == 1` is an invalid configuration, to use pure DP, use + `DistributedDataParallelKwargs` instead. + + """ + + dp_replicate_size: Optional[int] = None + dp_shard_size: Optional[int] = None + tp_size: Optional[int] = None + cp_size: Optional[int] = None + cp_backend: Literal["torch"] = None + sp_size: Optional[int] = None + sp_backend: Literal["deepspeed"] = None + + # we use Union because we might support other x parallel plugins (i.e. deepspeed, etc) + tp_handler: Union[None, TorchTensorParallelConfig] = None + cp_handler: Union[None, TorchContextParallelConfig] = None + sp_handler: Union[None, DeepSpeedSequenceParallelConfig] = None + + device_mesh = None + + def __repr__(self): + return ( + "ParallelismConfig(\n " + f"\tdp_replicate_size={self.dp_replicate_size},\n" + f"\tdp_shard_size={self.dp_shard_size},\n" + f"\ttp_size={self.tp_size},\n" + f"\tcp_size={self.cp_size},\n" + f"\tcp_backend={self.cp_backend},\n" + f"\tsp_size={self.sp_size},\n" + f"\tsp_backend={self.sp_backend},\n" + f"\ttotal_size={self.total_size}\n" + f"\ttp_handler={self.tp_handler},\n" + f"\tcp_handler={self.cp_handler})\n" + ) + + def to_json(self): + import copy + + _non_serializable_fields = ["device_mesh"] + + copy.deepcopy( + { + k: copy.deepcopy(v.__dict__) if hasattr(v, "__dict__") else v + for k, v in self.__dict__.items() + if k not in _non_serializable_fields + } + ) + + @property + def dp_dim_names(self): + """Names of enabled dimensions across which data parallelism is applied.""" + dims = [] + if self.dp_replicate_enabled: + dims += ["dp_replicate"] + if self.dp_shard_enabled: + dims += ["dp_shard"] + return dims + + @property + def non_dp_dim_names(self): + """Names of enabled dimensions which will receive the same batch (non-data parallel dimensions).""" + dims = [] + if self.tp_enabled: + dims += ["tp"] + if self.cp_enabled: + dims += ["cp"] + if self.sp_enabled: + dims += ["sp"] + return dims + + @property + def dp_shard_cp_dim_names(self): + """Names of enabled dimensions which will be flattened into a joint mesh across which is model sharded in FSDP.""" + dims = [] + if self.dp_shard_enabled: + dims += ["dp_shard"] + if self.cp_enabled: + dims += ["cp"] + return dims + + @property + def dp_cp_dim_names(self): + """Names of enabled dimensions across which loss should be averaged""" + dims = [] + if self.dp_replicate_enabled: + dims += ["dp_replicate"] + if self.dp_shard_enabled: + dims += ["dp_shard"] + if self.cp_enabled: + dims += ["cp"] + return dims + + @property + def fsdp_dim_names(self): + """Names of enabled dimensions across which FSDP is applied, including data parallel replication.""" + dims = [] + if self.dp_replicate_enabled: + dims += ["dp_replicate"] + dims += ["dp_shard_cp"] + return dims + + @property + def total_size(self): + """The total size of the parallelism configuration, which is the product of all sizes.""" + return self.dp_replicate_size * self.dp_shard_size * self.tp_size * self.cp_size * self.sp_size + + @property + def non_data_parallel_size(self): + """The size of the non-data parallel dimensions, which is the product of tensor and context parallel sizes.""" + return self.tp_size * self.cp_size * self.sp_size + + @property + def data_parallel_size(self): + """The size of the data parallel dimensions, which is the product of data parallel replication and""" + return self.dp_replicate_size * self.dp_shard_size + + @property + def dp_replicate_enabled(self): + """True if data parallel replication is enabled, i.e. `dp_replicate_size > 1`.""" + return self.dp_replicate_size > 1 + + @property + def dp_shard_enabled(self): + """True if data parallel sharding is enabled, i.e. `dp_shard_size > 1`.""" + return self.dp_shard_size > 1 + + @property + def tp_enabled(self): + """True if tensor parallelism is enabled, i.e. `tp_size > 1`.""" + return self.tp_size > 1 + + @property + def cp_enabled(self): + """True if context parallelism is enabled, i.e. `cp_size > 1`.""" + return self.cp_size > 1 + + @property + def sp_enabled(self): + """True if context parallelism is enabled, i.e. `sp_size > 1`.""" + return self.sp_size > 1 + + @property + def active_mesh_dims(self): + """Names of all active mesh dimensions.""" + return self.dp_dim_names + self.non_dp_dim_names + + def build_device_mesh(self, device_type: str): + """Builds a device mesh for the given device type based on the parallelism configuration. + This method will also create required joint meshes (e.g. `dp_shard_cp`, `dp_cp`, `dp`). + + Args: + device_type (`str`): The type of device for which to build the mesh, e + """ + if is_torch_version(">=", "2.2.0"): + from torch.distributed.device_mesh import init_device_mesh + else: + raise RuntimeError("Building a device_mesh requires to have torch>=2.2.0") + + mesh = self._get_mesh() + if len(mesh) == 0: + return None + mesh_dim_names, mesh_shape = mesh + device_mesh = init_device_mesh( + device_type, + mesh_shape, + mesh_dim_names=mesh_dim_names, + ) + if self.dp_dim_names: + device_mesh[self.dp_dim_names]._flatten("dp") + if self.dp_shard_cp_dim_names: + device_mesh[self.dp_shard_cp_dim_names]._flatten("dp_shard_cp") + if self.dp_cp_dim_names: + device_mesh[self.dp_cp_dim_names]._flatten("dp_cp") + + return device_mesh + + def get_device_mesh(self, device_type: Optional[str] = None): + if self.device_mesh is None: + if device_type is not None: + self.device_mesh = self.build_device_mesh(device_type) + else: + raise ("You need to pass a device_type e.g cuda to build the device mesh") + else: + if device_type is not None: + if self.device_mesh.device_type != device_type: + raise ValueError( + f"The device_mesh is already created with device type {self.device_mesh.device_type}. However, you are trying to get a device mesh with device_type {device_type}. Please check if you correctly initialized your device_mesh" + ) + return self.device_mesh + + def _get_mesh(self) -> tuple[tuple[int, ...], tuple[str, ...]]: + """Generate mesh shape and dimension names for torch.distributed.init_device_mesh().""" + + # Build mesh dimensions dictionary + mesh_dims = {parallelism: self._sizes[parallelism] for parallelism in self.active_mesh_dims} + + # Apply canonical ordering + mesh_order = ["dp_replicate", "dp_shard", "cp", "sp", "tp"] + sorted_items = sorted( + mesh_dims.items(), + key=lambda x: (mesh_order.index(x[0])), + ) + return tuple(zip(*sorted_items)) + + def __post_init__(self): + # Basic size validation + if self.dp_replicate_size is None: + self.dp_replicate_size = int(os.environ.get("PARALLELISM_CONFIG_DP_REPLICATE_SIZE", "1")) + if self.dp_shard_size is None: + self.dp_shard_size = int(os.environ.get("PARALLELISM_CONFIG_DP_SHARD_SIZE", "1")) + if self.tp_size is None: + self.tp_size = int(os.environ.get("PARALLELISM_CONFIG_TP_SIZE", "1")) + if self.cp_size is None: + self.cp_size = int(os.environ.get("PARALLELISM_CONFIG_CP_SIZE", "1")) + if self.cp_backend is None: + self.cp_backend = os.environ.get("PARALLELISM_CONFIG_CP_BACKEND", "torch") + if self.sp_size is None: + self.sp_size = int(os.environ.get("PARALLELISM_CONFIG_SP_SIZE", "1")) + if self.sp_backend is None: + self.sp_backend = os.environ.get("PARALLELISM_CONFIG_SP_BACKEND", "deepspeed") + + if self.tp_size > 1: + if self.tp_handler is None: + self.tp_handler = TorchTensorParallelConfig() + + if self.cp_size > 1: + if self.cp_handler is None: + self.cp_handler = TorchContextParallelConfig() + else: + cp_backends_config_map = dict( + torch=TorchContextParallelConfig, + ) + if not isinstance(self.cp_handler, cp_backends_config_map[self.cp_backend]): + raise ValueError( + f"ParallelismConfig's cp_backend={self.cp_backend} requires {cp_backends_config_map[self.cp_backend]}, but cp_handler was set to {type(self.cp_handler)}" + ) + + if self.sp_size > 1: + if self.sp_handler is None: + self.sp_handler = DeepSpeedSequenceParallelConfig() + if self.dp_replicate_size < 1: + raise ValueError(f"dp_replicate_size must be at least 1, but got {self.dp_replicate_size}") + if self.dp_shard_size < 1: + raise ValueError(f"dp_shard_size must be at least 1, but got {self.dp_shard_size}") + if self.tp_size < 1: + raise ValueError(f"tp_size must be at least 1, but got {self.tp_size}") + if self.cp_size < 1: + raise ValueError(f"cp_size must be at least 1, but got {self.cp_size}") + valid_cp_backends = ["torch"] + if self.cp_backend not in valid_cp_backends: + raise ValueError(f"cp_backend must be one of {valid_cp_backends}, but got {self.cp_backend}") + + if self.sp_size < 1: + raise ValueError(f"sp_size must be at least 1, but got {self.sp_size}") + valid_sp_backends = ["deepspeed"] + if self.sp_backend not in valid_sp_backends: + raise ValueError(f"sp_backend must be one of {valid_sp_backends}, but got {self.sp_backend}") + + if (self.tp_size > 1 or self.cp_size > 1) and self.dp_replicate_size > 1 and self.dp_shard_size == 1: + raise ValueError( + "Tensor/Context parallelism (tp/cp_size > 1) cannot be used with pure data parallelism (dp_replicate_size > 1 and dp_shard_size == 1). " + "Please set dp_shard_size > 1 and dp_replicate_size == 1 to compose FSDP + TP/CP for 2D parallel, " + "or set dp_replicate_size == 1 and dp_shard_size > 1 to compose HSDP + TP/CP for 3D parallel." + ) + self._sizes = { + "dp_replicate": self.dp_replicate_size, + "dp_shard": self.dp_shard_size, + "tp": self.tp_size, + "cp": self.cp_size, + "sp": self.sp_size, + } + + def _set_size(self, parallelism: str, size: int): + assert parallelism in self._sizes.keys(), f"Parallelism must be one of {self._sizes.keys()}" + self._sizes[parallelism] = size + setattr(self, f"{parallelism}_size", size) + + def _validate_accelerator(self, accelerator: "Accelerator"): + _warnings = set() + if not accelerator.multi_device and self.total_size == 1: + # No distributed setup, valid parallelism config + return + + # We need this to ensure DDP works + if self.total_size == 1: + self._set_size("dp_replicate", accelerator.num_processes) + + if self.total_size != accelerator.num_processes: + raise ValueError( + f"ParallelismConfig total_size ({self.total_size}) does not match " + f"num_processes ({accelerator.num_processes}). Please adjust dp_replicate_size/ " + f"dp_shard_size/tp_size/cp_size/sp_size." + ) + + if self.total_size > 1 and not ( + accelerator.is_fsdp2 + or accelerator.multi_device + or accelerator.distributed_type == DistributedType.DEEPSPEED + ): + raise ValueError( + f"ParallelismConfig is only compatible DistributedType.FSDP (version 2) or DistributedType.Multi{{Device}} or DistributedType.DEEPSPEED, but got {accelerator.distributed_type}." + ) + + for parallelism, size in self._sizes.items(): + if size == 1 and getattr(self, f"{parallelism}_handler", None) is not None: + _warnings.add( + f"ParallelismConfig.{parallelism}_handler is set, but {parallelism}_size is set to 1. This handler will be ignored." + ) + + if _warnings and accelerator.is_main_process: + warnings.warn( + "ParallelismConfig has the following warnings:\n" + "\n".join(_warnings), + UserWarning, + ) diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/scheduler.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/scheduler.py new file mode 100644 index 0000000000000000000000000000000000000000..1fa8a13f238afd7b908ee8e8cb8e0620f48d4ff8 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/scheduler.py @@ -0,0 +1,98 @@ +# Copyright 2022 The HuggingFace Team. All rights reserved. +# +# 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. + +# We ignore warnings about stepping the scheduler since we step it ourselves during gradient accumulation + +import warnings + +from .state import AcceleratorState, GradientState + + +warnings.filterwarnings("ignore", category=UserWarning, module="torch.optim.lr_scheduler") + + +class AcceleratedScheduler: + """ + A wrapper around a learning rate scheduler that will only step when the optimizer(s) have a training step. Useful + to avoid making a scheduler step too fast when gradients went overflow and there was no training step (in mixed + precision training) + + When performing gradient accumulation scheduler lengths should not be changed accordingly, Accelerate will always + step the scheduler to account for it. + + Args: + scheduler (`torch.optim.lr_scheduler._LRScheduler`): + The scheduler to wrap. + optimizers (one or a list of `torch.optim.Optimizer`): + The optimizers used. + step_with_optimizer (`bool`, *optional*, defaults to `True`): + Whether or not the scheduler should be stepped at each optimizer step. + split_batches (`bool`, *optional*, defaults to `False`): + Whether or not the dataloaders split one batch across the different processes (so batch size is the same + regardless of the number of processes) or create batches on each process (so batch size is the original + batch size multiplied by the number of processes). + """ + + def __init__(self, scheduler, optimizers, step_with_optimizer: bool = True, split_batches: bool = False): + self.scheduler = scheduler + self.optimizers = optimizers if isinstance(optimizers, (list, tuple)) else [optimizers] + self.split_batches = split_batches + self.step_with_optimizer = step_with_optimizer + self.gradient_state = GradientState() + + def step(self, *args, **kwargs): + if not self.step_with_optimizer: + # No link between scheduler and optimizer -> just step + self.scheduler.step(*args, **kwargs) + return + + # Otherwise, first make sure the optimizer was stepped. + if not self.gradient_state.sync_gradients: + if self.gradient_state.adjust_scheduler: + self.scheduler._step_count += 1 + return + + for opt in self.optimizers: + if opt.step_was_skipped: + return + if self.split_batches: + # Split batches -> the training dataloader batch size is not changed so one step per training step + self.scheduler.step(*args, **kwargs) + else: + # Otherwise the training dataloader batch size was multiplied by `num_processes`, so we need to do + # num_processes steps per training step + num_processes = AcceleratorState().num_processes + for _ in range(num_processes): + # Special case when using OneCycle and `drop_last` was not used + if hasattr(self.scheduler, "total_steps"): + if self.scheduler._step_count <= self.scheduler.total_steps: + self.scheduler.step(*args, **kwargs) + else: + self.scheduler.step(*args, **kwargs) + + # Passthroughs + def get_last_lr(self): + return self.scheduler.get_last_lr() + + def state_dict(self): + return self.scheduler.state_dict() + + def load_state_dict(self, state_dict): + self.scheduler.load_state_dict(state_dict) + + def get_lr(self): + return self.scheduler.get_lr() + + def print_lr(self, *args, **kwargs): + return self.scheduler.print_lr(*args, **kwargs) diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/state.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/state.py new file mode 100644 index 0000000000000000000000000000000000000000..4da3c91f3c29275eb109d5116b7dd07bd88a29f6 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/state.py @@ -0,0 +1,1373 @@ +# Copyright 2021 The HuggingFace Team. All rights reserved. +# +# 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. + +from __future__ import annotations + +import logging +import os +import threading +import warnings +import weakref +from contextlib import contextmanager +from functools import partial +from typing import Any, Callable + +import torch + +from .utils import ( + DistributedType, + DynamoBackend, + GradientAccumulationPlugin, + check_cuda_fp8_capability, + check_cuda_p2p_ib_support, + deepspeed_required, + get_cpu_distributed_information, + get_int_from_env, + is_ccl_available, + is_datasets_available, + is_deepspeed_available, + is_fp8_available, + is_habana_gaudi1, + is_hpu_available, + is_ipex_available, + is_mlu_available, + is_mps_available, + is_musa_available, + is_npu_available, + is_sdaa_available, + is_torch_xla_available, + is_xccl_available, + is_xpu_available, + parse_choice_from_env, + parse_flag_from_env, + set_numa_affinity, +) +from .utils.dataclasses import SageMakerDistributedType + + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + import torch_xla.runtime as xr + +if is_mlu_available(check_device=False): + import torch_mlu # noqa: F401 + +if is_sdaa_available(check_device=False): + import torch_sdaa # noqa: F401 + +if is_musa_available(check_device=False): + import torch_musa # noqa: F401 + +if is_npu_available(check_device=False): + import torch_npu # noqa: F401 + + +logger = logging.getLogger(__name__) + + +def is_initialized() -> bool: + """ + Checks if the `AcceleratorState` has been initialized from `Accelerator`. Same as `AcceleratorState.initialized`, + but works as a module method. + """ + return AcceleratorState._shared_state != {} + + +# Lambda function that does nothing +def do_nothing(*args, **kwargs): + return None + + +class ThreadLocalSharedDict(threading.local): + """ + Descriptor that holds a dict shared between instances of a class in the same thread. + + Note: Descriptors have slightly different semantics than just a dict field on its own. + `PartialState(...)._shared_state` and `PartialState._shared_state` (instance vs class) give the same value: the + underlying _storage dict. Likewise, `PartialState(...)._shared_state = {...}` overrides the _storage dict inside + the descriptor as you would expect. However, `PartialState._shared_state = {}` actually replaces the descriptor + object with a dict instead Thus, you should modify the _storage dict in-place (e.g. `_shared_state.clear()`). + + See Python documentation for an explanation of descriptors: https://docs.python.org/3/howto/descriptor.html + + This is required for using PyTorch/XLA with PJRT in multithreaded mode (required for TPU v2 and v3). + + See https://github.com/pytorch/xla/blob/r2.0/docs/pjrt.md#multithreading-on-tpu-v2v3 + """ + + def __init__(self, thread_local: bool = False): + self._storage = {} + + def __get__(self, obj, objtype=None): + return self._storage + + def __set__(self, obj, value): + self._storage = value + + +# Prefer global shared dictionary, except when using TPU. +SharedDict = dict if not is_torch_xla_available() else ThreadLocalSharedDict + + +# Inspired by Alex Martelli's 'Borg'. +class PartialState: + """ + Singleton class that has information about the current training environment and functions to help with process + control. Designed to be used when only process control and device execution states are needed. Does *not* need to + be initialized from `Accelerator`. + + Args: + cpu (`bool`, *optional*): + Whether or not to force the script to execute on CPU. Will ignore any accelerators available if set to + `True` and force the execution on the CPU. + kwargs (additional keyword arguments, *optional*): + Additional keyword arguments to pass to the relevant `init_process_group` function. Valid `kwargs` can be + found in [`utils.InitProcessGroupKwargs`]. See the example section for detailed usage. + + **Available attributes:** + + - **device** (`torch.device`) -- The device to use. + - **distributed_type** ([`~accelerate.state.DistributedType`]) -- The type of distributed environment currently + in use. + - **local_process_index** (`int`) -- The index of the current process on the current server. + - **mixed_precision** (`str`) -- Whether or not the current script will use mixed precision, and if so the type + of mixed precision being performed. (Choose from 'no','fp16','bf16 or 'fp8'). + - **num_processes** (`int`) -- The number of processes currently launched in parallel. + - **process_index** (`int`) -- The index of the current process. + - **is_last_process** (`bool`) -- Whether or not the current process is the last one. + - **is_main_process** (`bool`) -- Whether or not the current process is the main one. + - **is_local_main_process** (`bool`) -- Whether or not the current process is the main one on the local node. + - **debug** (`bool`) -- Whether or not the current script is being run in debug mode. + + Example: + ```python + from accelerate.utils import InitProcessGroupKwargs + + # To include `InitProcessGroupKwargs`, init then call `.to_kwargs()` + kwargs = InitProcessGroupKwargs(...).to_kwargs() + state = PartialState(**kwargs) + ``` + """ + + _shared_state = SharedDict() + _known_attrs = [ + "_cpu", + "_mixed_precision", + "_shared_state", + "backend", + "debug", + "device", + "distributed_type", + "fork_launched", + "local_process_index", + "num_processes", + "process_index", + ] + + def __init__(self, cpu: bool = False, **kwargs): + self.__dict__ = self._shared_state + if not self.initialized: + self._cpu = cpu + self.backend = None + env_device = os.environ.get("ACCELERATE_TORCH_DEVICE", None) + self.device = torch.device(env_device) if env_device is not None else None + self.debug = parse_flag_from_env("ACCELERATE_DEBUG_MODE") + use_sagemaker_dp = kwargs.pop("_use_sagemaker_dp", None) + dist_information = None + if use_sagemaker_dp is None: + use_sagemaker_dp = ( + os.environ.get("ACCELERATE_USE_SAGEMAKER", "false").lower() == "true" + and os.environ.get("ACCELERATE_SAGEMAKER_DISTRIBUTED_TYPE") != SageMakerDistributedType.NO + ) + + # Sets up self.backend + imports + original_backend = kwargs.pop("backend", None) + backend, distributed_type = self._prepare_backend(cpu, use_sagemaker_dp, original_backend) + if original_backend is not None and backend != original_backend: + raise ValueError(f"Your assigned backend {original_backend} is not available, please use {backend}") + self.backend = backend + self.distributed_type = distributed_type + use_deepspeed = False + if not cpu and self.backend != "xla": + if int(os.environ.get("LOCAL_RANK", -1)) != -1: + # Deal with spawning deepspeed + if os.environ.get("ACCELERATE_USE_DEEPSPEED", "false").lower() == "true": + if not is_deepspeed_available(): + raise ImportError( + "DeepSpeed is not available => install it using `pip3 install deepspeed` or build it from source" + ) + from deepspeed import comm as dist + + if not dist.is_initialized(): + if self.backend == "tccl": + local_rank = os.environ.get("LOCAL_RANK", -1) + torch.sdaa.set_device(f"sdaa:{local_rank}") + dist.init_distributed(dist_backend=self.backend, auto_mpi_discovery=False, **kwargs) + # We need to flag to `use_deepspeed` to be True to override `distributed_type` later + use_deepspeed = True + # Deal with all other backends but XPU and CPU, that gets handled special later + elif ( + self.distributed_type not in (DistributedType.MULTI_XPU, DistributedType.MULTI_CPU) + and not torch.distributed.is_initialized() + ): + if self.backend == "tccl": + local_rank = os.environ.get("LOCAL_RANK", -1) + torch.sdaa.set_device(f"sdaa:{local_rank}") + if ( + self.backend == "nccl" + and os.environ.get("ACCELERATE_USE_FSDP", "false").lower() == "true" + and ( + os.environ.get("FSDP_OFFLOAD_PARAMS", "false").lower() == "true" + or os.environ.get("FSDP_STATE_DICT_TYPE", "SHARDED_STATE_DICT") == "FULL_STATE_DICT" + ) + ): + self.backend = "cuda:nccl,cpu:gloo" + torch.distributed.init_process_group(backend=self.backend, **kwargs) + + # XPU and CPU require special env configs to be set + if self.distributed_type in (DistributedType.MULTI_XPU, DistributedType.MULTI_CPU): + dist_information = get_cpu_distributed_information() + os.environ["RANK"] = str(dist_information.rank) + os.environ["WORLD_SIZE"] = str(dist_information.world_size) + os.environ["LOCAL_RANK"] = str(dist_information.local_rank) + os.environ["LOCAL_WORLD_SIZE"] = str(dist_information.local_world_size) + if not os.environ.get("MASTER_PORT", None): + os.environ["MASTER_PORT"] = "29500" + if ( + not os.environ.get("MASTER_ADDR", None) + and dist_information.local_world_size != dist_information.world_size + and self.backend != "mpi" + ): + raise ValueError( + "Tried to launch on distributed with multinode, but `MASTER_ADDR` env was not set, " + "please try exporting rank 0's hostname as `MASTER_ADDR`" + ) + kwargs["rank"] = dist_information.rank + kwargs["world_size"] = dist_information.world_size + + if ( + self.distributed_type == DistributedType.MULTI_CPU + and get_int_from_env(["OMP_NUM_THREADS"], 0) == 0 + ): + import psutil + + num_cpu_threads_per_process = int( + psutil.cpu_count(logical=False) / dist_information.local_world_size + ) + if num_cpu_threads_per_process == 0: + num_cpu_threads_per_process = 1 + torch.set_num_threads(num_cpu_threads_per_process) + warnings.warn( + f"OMP_NUM_THREADS/MKL_NUM_THREADS unset, we set it at {num_cpu_threads_per_process} to improve oob" + " performance." + ) + + if not torch.distributed.is_initialized(): + torch.distributed.init_process_group(backend=self.backend, **kwargs) + + # No backend == no distributed training + if self.backend is None: + self.distributed_type = DistributedType.NO + self.num_processes = 1 + self.process_index = 0 + self.local_process_index = 0 + elif self.backend == "xla": + # XLA needs device setting first for `set_replication` + self.set_device() + xm.set_replication(self.device, xm.get_xla_supported_devices()) + self.num_processes = xr.world_size() + self.process_index = xr.global_ordinal() + if is_torch_xla_available(check_is_tpu=True): + self.local_process_index = xm.get_local_ordinal() + else: + self.local_process_index = int(os.environ.get("LOCAL_RANK", -1)) + else: + self.num_processes = torch.distributed.get_world_size() + self.process_index = torch.distributed.get_rank() + self.local_process_index = ( + int(os.environ.get("LOCAL_RANK", -1)) if dist_information is None else dist_information.local_rank + ) + self.set_device() + # Now we can change to deepseed + if use_deepspeed: + self.distributed_type = DistributedType.DEEPSPEED + + # Set CPU affinity if enabled + if parse_flag_from_env("ACCELERATE_CPU_AFFINITY", False): + set_numa_affinity(self.local_process_index) + + # Check for old RTX 4000's that can't use P2P or IB and are on old drivers + if self.device.type == "cuda" and not check_cuda_p2p_ib_support(): + if "NCCL_P2P_DISABLE" not in os.environ or "NCCL_IB_DISABLE" not in os.environ: + raise NotImplementedError( + "Using RTX 4000 series doesn't support faster communication broadband via P2P or IB. " + 'Please set `NCCL_P2P_DISABLE="1"` and `NCCL_IB_DISABLE="1" or use `accelerate launch` which ' + "will do this automatically." + ) + + # Important: This should be the *only* code outside of `self.initialized!` + self.fork_launched = parse_flag_from_env("FORK_LAUNCHED", 0) + + def __repr__(self) -> str: + return ( + f"Distributed environment: {self.distributed_type}{(' Backend: ' + self.backend) if self.backend else ''}\n" + f"Num processes: {self.num_processes}\n" + f"Process index: {self.process_index}\n" + f"Local process index: {self.local_process_index}\n" + f"Device: {self.device}\n" + ) + + @staticmethod + def _reset_state(): + "Resets `_shared_state`, is used internally and should not be called" + PartialState._shared_state.clear() + + @property + def initialized(self) -> bool: + "Returns whether the `PartialState` has been initialized" + return self._shared_state != {} + + @property + def use_distributed(self): + """ + Whether the Accelerator is configured for distributed training + """ + return self.distributed_type != DistributedType.NO and self.num_processes > 1 + + @property + def is_last_process(self) -> bool: + "Returns whether the current process is the last one" + return self.process_index == self.num_processes - 1 + + @property + def is_main_process(self) -> bool: + "Returns whether the current process is the main process" + return ( + self.process_index == 0 if self.distributed_type != DistributedType.MEGATRON_LM else self.is_last_process + ) + + @property + def is_local_main_process(self) -> bool: + "Returns whether the current process is the main process on the local node" + return ( + self.local_process_index == 0 + if self.distributed_type != DistributedType.MEGATRON_LM + else self.is_last_process + ) + + def wait_for_everyone(self): + """ + Will stop the execution of the current process until every other process has reached that point (so this does + nothing when the script is only run in one process). Useful to do before saving a model. + + Example: + + ```python + >>> # Assuming two GPU processes + >>> import time + >>> from accelerate.state import PartialState + + >>> state = PartialState() + >>> if state.is_main_process: + ... time.sleep(2) + >>> else: + ... print("I'm waiting for the main process to finish its sleep...") + >>> state.wait_for_everyone() + >>> # Should print on every process at the same time + >>> print("Everyone is here") + ``` + """ + if self.distributed_type in ( + DistributedType.MULTI_GPU, + DistributedType.MULTI_MLU, + DistributedType.MULTI_SDAA, + DistributedType.MULTI_MUSA, + DistributedType.MULTI_NPU, + DistributedType.MULTI_XPU, + DistributedType.MULTI_CPU, + DistributedType.MULTI_HPU, + DistributedType.DEEPSPEED, + DistributedType.FSDP, + ): + torch.distributed.barrier(device_ids=[self.local_process_index]) + elif self.distributed_type == DistributedType.XLA: + xm.rendezvous("accelerate.utils.wait_for_everyone") + + def _goes_first(self, is_main: bool): + if not is_main: + self.wait_for_everyone() + + yield + + if is_main: + self.wait_for_everyone() + + @contextmanager + def split_between_processes(self, inputs: list | tuple | dict | torch.Tensor, apply_padding: bool = False): + """ + Splits `input` between `self.num_processes` quickly and can be then used on that process. Useful when doing + distributed inference, such as with different prompts. + + Note that when using a `dict`, all keys need to have the same number of elements. + + Args: + inputs (`list`, `tuple`, `torch.Tensor`, `dict` of `list`/`tuple`/`torch.Tensor`, or `datasets.Dataset`): + The input to split between processes. + apply_padding (`bool`, `optional`, defaults to `False`): + Whether to apply padding by repeating the last element of the input so that all processes have the same + number of elements. Useful when trying to perform actions such as `gather()` on the outputs or passing + in less inputs than there are processes. If so, just remember to drop the padded elements afterwards. + + + Example: + + ```python + # Assume there are two processes + from accelerate import PartialState + + state = PartialState() + with state.split_between_processes(["A", "B", "C"]) as inputs: + print(inputs) + # Process 0 + ["A", "B"] + # Process 1 + ["C"] + + with state.split_between_processes(["A", "B", "C"], apply_padding=True) as inputs: + print(inputs) + # Process 0 + ["A", "B"] + # Process 1 + ["C", "C"] + ``` + """ + if self.num_processes == 1: + yield inputs + return + length = len(inputs) + # Nested dictionary of any types + if isinstance(inputs, dict): + length = len(inputs[list(inputs.keys())[0]]) + if not all(len(v) == length for v in inputs.values()): + raise ValueError("All values in the dictionary must have the same length") + num_samples_per_process, num_extras = divmod(length, self.num_processes) + start_index = self.process_index * num_samples_per_process + min(self.process_index, num_extras) + end_index = start_index + num_samples_per_process + (1 if self.process_index < num_extras else 0) + + def _split_values(inputs, start_index, end_index): + if isinstance(inputs, (list, tuple, torch.Tensor)): + if start_index >= len(inputs): + result = inputs[-1:] + else: + result = inputs[start_index:end_index] + if apply_padding: + if isinstance(result, torch.Tensor): + from accelerate.utils import pad_across_processes, send_to_device + + # The tensor needs to be on the device before we can pad it + tensorized_result = send_to_device(result, self.device) + result = pad_across_processes(tensorized_result, pad_index=inputs[-1]) + else: + result += [result[-1]] * (num_samples_per_process + (1 if num_extras > 0 else 0) - len(result)) + return result + elif isinstance(inputs, dict): + for key in inputs.keys(): + inputs[key] = _split_values(inputs[key], start_index, end_index) + return inputs + else: + if is_datasets_available(): + from datasets import Dataset + + if isinstance(inputs, Dataset): + if start_index >= len(inputs): + start_index = len(inputs) - 1 + if end_index > len(inputs): + end_index = len(inputs) + result_idcs = list(range(start_index, end_index)) + if apply_padding: + result_idcs += [end_index - 1] * ( + num_samples_per_process + (1 if num_extras > 0 else 0) - len(result_idcs) + ) + return inputs.select(result_idcs) + return inputs + + yield _split_values(inputs, start_index, end_index) + + @contextmanager + def main_process_first(self): + """ + Lets the main process go first inside a with block. + + The other processes will enter the with block after the main process exits. + + Example: + + ```python + >>> from accelerate import Accelerator + + >>> accelerator = Accelerator() + >>> with accelerator.main_process_first(): + ... # This will be printed first by process 0 then in a seemingly + ... # random order by the other processes. + ... print(f"This will be printed by process {accelerator.process_index}") + ``` + """ + yield from self._goes_first(self.is_main_process) + + @contextmanager + def local_main_process_first(self): + """ + Lets the local main process go inside a with block. + + The other processes will enter the with block after the main process exits. + + Example: + + ```python + >>> from accelerate.state import PartialState + + >>> state = PartialState() + >>> with state.local_main_process_first(): + ... # This will be printed first by local process 0 then in a seemingly + ... # random order by the other processes. + ... print(f"This will be printed by process {state.local_process_index}") + ``` + """ + yield from self._goes_first(self.is_local_main_process) + + def on_main_process(self, function: Callable[..., Any] | None = None): + """ + Decorator that only runs the decorated function on the main process. + + Args: + function (`Callable`): The function to decorate. + + Example: + + ```python + >>> from accelerate.state import PartialState + + >>> state = PartialState() + + + >>> @state.on_main_process + ... def print_something(): + ... print("This will be printed by process 0 only.") + + + >>> print_something() + "This will be printed by process 0 only" + ``` + """ + if not self.initialized: + raise ValueError("The `PartialState` or `Accelerator` must be initialized before calling this function.") + if self.is_main_process or not self.use_distributed: + return function + return do_nothing + + def on_local_main_process(self, function: Callable[..., Any] | None = None): + """ + Decorator that only runs the decorated function on the local main process. + + Args: + function (`Callable`): The function to decorate. + + Example: + ```python + # Assume we have 2 servers with 4 processes each. + from accelerate.state import PartialState + + state = PartialState() + + + @state.on_local_main_process + def print_something(): + print("This will be printed by process 0 only on each server.") + + + print_something() + # On server 1: + "This will be printed by process 0 only" + # On server 2: + "This will be printed by process 0 only" + ``` + """ + if self.is_local_main_process or not self.use_distributed: + return function + return do_nothing + + def on_last_process(self, function: Callable[..., Any]): + """ + Decorator that only runs the decorated function on the last process. + + Args: + function (`Callable`): The function to decorate. + + Example: + ```python + # Assume we have 4 processes. + from accelerate.state import PartialState + + state = PartialState() + + + @state.on_last_process + def print_something(): + print(f"Printed on process {state.process_index}") + + + print_something() + "Printed on process 3" + ``` + """ + if self.is_last_process or not self.use_distributed: + return function + return do_nothing + + def on_process(self, function: Callable[..., Any] | None = None, process_index: int | None = None): + """ + Decorator that only runs the decorated function on the process with the given index. + + Args: + function (`Callable`, `optional`): + The function to decorate. + process_index (`int`, `optional`): + The index of the process on which to run the function. + + Example: + ```python + # Assume we have 4 processes. + from accelerate.state import PartialState + + state = PartialState() + + + @state.on_process(process_index=2) + def print_something(): + print(f"Printed on process {state.process_index}") + + + print_something() + "Printed on process 2" + ``` + """ + if function is None: + return partial(self.on_process, process_index=process_index) + if (self.process_index == process_index) or (not self.use_distributed): + return function + return do_nothing + + def on_local_process(self, function: Callable[..., Any] | None = None, local_process_index: int | None = None): + """ + Decorator that only runs the decorated function on the process with the given index on the current node. + + Args: + function (`Callable`, *optional*): + The function to decorate. + local_process_index (`int`, *optional*): + The index of the local process on which to run the function. + + Example: + ```python + # Assume we have 2 servers with 4 processes each. + from accelerate import Accelerator + + accelerator = Accelerator() + + + @accelerator.on_local_process(local_process_index=2) + def print_something(): + print(f"Printed on process {accelerator.local_process_index}") + + + print_something() + # On server 1: + "Printed on process 2" + # On server 2: + "Printed on process 2" + ``` + """ + if function is None: + return partial(self.on_local_process, local_process_index=local_process_index) + if (self.local_process_index == local_process_index) or (not self.use_distributed): + return function + return do_nothing + + def print(self, *args, **kwargs): + if self.is_local_main_process: + print(*args, **kwargs) + + @property + def default_device(self) -> torch.device: + """ + Returns the default device which is: + - MPS if `torch.backends.mps.is_available()` and `torch.backends.mps.is_built()` both return True. + - CUDA if `torch.cuda.is_available()` + - MLU if `is_mlu_available()` + - SDAA if `is_sdaa_available()` + - MUSA if `is_musa_available()` + - NPU if `is_npu_available()` + - HPU if `is_hpu_available()` + - CPU otherwise + """ + if is_mps_available(): + os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" + return torch.device("mps") + elif is_mlu_available(): + return torch.device("mlu") + elif is_sdaa_available(): + return torch.device("sdaa") + elif is_musa_available(): + return torch.device("musa") + # NPU should be checked before CUDA when using `transfer_to_npu` + # See issue #3020: https://github.com/huggingface/accelerate/issues/3020 + elif is_npu_available(): + return torch.device("npu") + elif is_hpu_available(): + return torch.device("hpu") + elif torch.cuda.is_available(): + return torch.device("cuda") + elif is_xpu_available(): + return torch.device("xpu") + else: + return torch.device("cpu") + + def _prepare_backend( + self, cpu: bool = False, sagemaker_dp=False, backend: str | None = None + ) -> tuple[str, DistributedType]: + "Prepares any imports needed before initializing the distributed backend and sets `self.backend` properly" + distributed_type = None + if sagemaker_dp: + import smdistributed.dataparallel.torch.torch_smddp # noqa + + backend = "smddp" + distributed_type = DistributedType.MULTI_GPU + elif is_torch_xla_available(): + backend = "xla" + distributed_type = DistributedType.XLA + + elif int(os.environ.get("LOCAL_RANK", -1)) != -1 and not cpu: + if is_mlu_available(): + backend = "cncl" + distributed_type = DistributedType.MULTI_MLU + if is_sdaa_available(): + backend = "tccl" + distributed_type = DistributedType.MULTI_SDAA + elif is_musa_available(): + backend = "mccl" + distributed_type = DistributedType.MULTI_MUSA + # NPU should be checked before CUDA when using `transfer_to_npu` + # See issue #3020: https://github.com/huggingface/accelerate/issues/3020 + elif is_npu_available(): + backend = "hccl" + distributed_type = DistributedType.MULTI_NPU + elif is_hpu_available(init_hccl=True): + if backend is None: + backend = "hccl" + distributed_type = DistributedType.MULTI_HPU + elif torch.cuda.is_available(): + if backend is None: + backend = "nccl" + distributed_type = DistributedType.MULTI_GPU + elif is_xpu_available() and is_xccl_available(): + if backend is None: + backend = "xccl" + distributed_type = DistributedType.MULTI_XPU + + if distributed_type is None and ( + int(os.environ.get("LOCAL_RANK", -1)) != -1 + or get_int_from_env(["PMI_SIZE", "OMPI_COMM_WORLD_SIZE", "MV2_COMM_WORLD_SIZE", "WORLD_SIZE"], 1) > 1 + ): + if not cpu and is_xpu_available(): + distributed_type = DistributedType.MULTI_XPU + else: + distributed_type = DistributedType.MULTI_CPU + + if ( + backend in (None, "ccl") + and is_ccl_available() + and (get_int_from_env(["CCL_WORKER_COUNT"], 0) > 0 or distributed_type == DistributedType.MULTI_XPU) + ): + import oneccl_bindings_for_pytorch # noqa: F401 + + backend = "ccl" + elif backend in (None, "mpi") and torch.distributed.is_mpi_available(): + backend = "mpi" + else: + backend = "gloo" + if distributed_type is None: + distributed_type = DistributedType.NO + + return backend, distributed_type + + def set_device(self): + """ + Sets the device in `self.device` to the current distributed environment. + """ + if self.device is not None: + return + if self.distributed_type == DistributedType.NO: + self.device = torch.device("cpu") if self._cpu else self.default_device + return + device = str(self.distributed_type).split(".")[-1].replace("MULTI_", "").lower() + if device not in ("cpu", "gpu", "mlu", "musa", "npu", "xpu", "xla", "hpu", "sdaa"): + raise ValueError( + f"Can't set device for {self.distributed_type} ({device}), verify we should be calling `_set_device()` for it!" + ) + if device == "xla": + self.device = xm.xla_device() + elif device == "hpu": + self.device = torch.device("hpu", torch.hpu.current_device()) + else: + if device == "gpu": + device = "cuda" + device_module = getattr(torch, device) + device_index = self.local_process_index % device_module.device_count() + self.device = torch.device(device, device_index) + device_module.set_device(self.device) + + def destroy_process_group(self, group=None): + """ + Destroys the process group. If one is not specified, the default process group is destroyed. + """ + if self.fork_launched and group is None: + return + # needed when using torch.distributed.init_process_group + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group(group) + + def __getattr__(self, name: str): + # By this point we know that no attributes of `self` contain `name`, + # so we just modify the error message + if name in self._known_attrs: + raise AttributeError( + f"`PartialState` object has no attribute `{name}`. " + "This happens if `PartialState._reset_state()` was called and " + "an `Accelerator` or `PartialState` was not reinitialized." + ) + # Raise a typical AttributeError + raise AttributeError(f"'PartialState' object has no attribute '{name}'") + + +class AcceleratorState: + """ + Singleton class that has information about the current training environment. + + **Available attributes:** + + - **device** (`torch.device`) -- The device to use. + - **distributed_type** ([`~accelerate.state.DistributedType`]) -- The type of distributed environment currently + in use. + - **parallelism_config** ([`~accelerate.utils.ParallelismConfig`]) -- The parallelism configuration for the + current training environment. This is used to configure the distributed training environment. + - **initialized** (`bool`) -- Whether or not the `AcceleratorState` has been initialized from `Accelerator`. + - **local_process_index** (`int`) -- The index of the current process on the current server. + - **mixed_precision** (`str`) -- Whether or not the current script will use mixed precision, and if so the type + of mixed precision being performed. (Choose from 'no','fp16','bf16 or 'fp8'). + - **num_processes** (`int`) -- The number of processes currently launched in parallel. + - **process_index** (`int`) -- The index of the current process. + - **is_last_process** (`bool`) -- Whether or not the current process is the last one. + - **is_main_process** (`bool`) -- Whether or not the current process is the main one. + - **is_local_main_process** (`bool`) -- Whether or not the current process is the main one on the local node. + - **debug** (`bool`) -- Whether or not the current script is being run in debug mode. + """ + + _shared_state = SharedDict() + _known_attrs = PartialState._known_attrs + [ + "deepspeed_plugin", + "use_ipex", + "fsdp_plugin", + "megatron_lm_plugin", + "dynamo_plugin", + ] + + def __init__( + self, + mixed_precision: str | None = None, + cpu: bool = False, + dynamo_plugin=None, + deepspeed_plugin=None, + fsdp_plugin=None, + torch_tp_plugin=None, + megatron_lm_plugin=None, + parallelism_config=None, + _from_accelerator: bool = False, + **kwargs, + ): + self.__dict__ = self._shared_state + if parse_flag_from_env("ACCELERATE_USE_CPU"): + cpu = True + if PartialState._shared_state == {}: + PartialState(cpu, **kwargs) + self.__dict__.update(PartialState._shared_state) + self._check_initialized(mixed_precision, cpu) + if not self.initialized: + self.deepspeed_plugins = None + self.use_ipex = None + self.torch_tp_plugin = torch_tp_plugin + self.parallelism_config = parallelism_config + self.device_mesh = None + mixed_precision = ( + parse_choice_from_env("ACCELERATE_MIXED_PRECISION", "no") + if mixed_precision is None + else mixed_precision.lower() + ) + if mixed_precision == "fp8": + # this is confusing, why is is_fp8_available only checks for library availability ? + if not is_fp8_available(): + raise ValueError( + "Using `fp8` precision requires `transformer_engine` or `MS-AMP` to be installed." + ) + elif torch.cuda.is_available() and not check_cuda_fp8_capability(): + logger.warning( + f"The current device has compute capability of {torch.cuda.get_device_capability()} which is " + "insufficient for FP8 mixed precision training (requires a GPU Hopper/Ada Lovelace " + "or higher, compute capability of 8.9 or higher). Will use FP16 instead." + ) + mixed_precision = "fp16" + elif is_habana_gaudi1(): + logger.warning( + "The current HPU device is Gaudi1 which does not support FP8 mixed precision training (requires " + "Gaudi2 or higher). Will use BF16 instead." + ) + mixed_precision = "bf16" + + self.dynamo_plugin = dynamo_plugin + if not _from_accelerator: + raise ValueError( + "Please make sure to properly initialize your accelerator via `accelerator = Accelerator()` " + "before using any functionality from the `accelerate` library." + ) + # deepspeed handles mixed_precision using deepspeed_config. But we need to set it to fp8 + # if we're using fp8. + if self.distributed_type == DistributedType.DEEPSPEED and mixed_precision != "fp8": + self._mixed_precision = "no" + else: + self._mixed_precision = mixed_precision + + if self.distributed_type == DistributedType.XLA and is_torch_xla_available(check_is_tpu=True): + if mixed_precision == "bf16": + if os.environ.get("ACCELERATE_DOWNCAST_BF16"): + os.environ["XLA_USE_BF16"] = str(0) + os.environ["XLA_DOWNCAST_BF16"] = str(1) + self.downcast_bfloat = True + else: + os.environ["XLA_USE_BF16"] = str(1) + os.environ["XLA_DOWNCAST_BF16"] = str(0) + self.downcast_bfloat = False + elif os.environ.get("ACCELERATE_USE_DEEPSPEED", "false").lower() == "true" and not cpu: + self.distributed_type = DistributedType.DEEPSPEED + if not isinstance(deepspeed_plugin, dict): + deepspeed_plugin.set_mixed_precision(mixed_precision) + deepspeed_plugin.select(_from_accelerator_state=True) + else: + for plugin in deepspeed_plugin.values(): + plugin.set_mixed_precision(mixed_precision) + # The first plugin passed in is always the active one + first_plugin = next(iter(deepspeed_plugin.values())) + first_plugin.select(_from_accelerator_state=True) + self.deepspeed_plugins = deepspeed_plugin + elif self.distributed_type in [ + DistributedType.MULTI_GPU, + DistributedType.MULTI_MLU, + DistributedType.MULTI_SDAA, + DistributedType.MULTI_MUSA, + DistributedType.MULTI_NPU, + DistributedType.MULTI_XPU, + DistributedType.MULTI_HPU, + ]: + # TODO: Siro - remove when axolotl fixes their side + if not os.environ.get("ACCELERATE_ALLOW_CP_STANDALONE", "false").lower() == "true": + if self.parallelism_config and self.parallelism_config.cp_enabled and fsdp_plugin is None: + raise ValueError( + "`cp_size > 1` specified in the `parallelism_config`, but no `fsdp_plugin` was provided. We need a `fsdp_plugin` to use context parallelism with `cp_backend=torch`, as we also shard the model across the device mesh to save more memory" + ) + if ( + self.parallelism_config is not None + and self.parallelism_config.cp_enabled + and fsdp_plugin.fsdp_version == 1 + ): + raise ValueError( + "Using `cp_size>1` requires FSDP2, but the provided `fsdp_plugin` is using FSDP1. " + ) + if (os.environ.get("ACCELERATE_USE_FSDP", "false").lower() == "true" or fsdp_plugin is not None) or ( + self.parallelism_config is not None and self.parallelism_config.cp_enabled + ): + self.distributed_type = DistributedType.FSDP + if self._mixed_precision != "no" and fsdp_plugin is not None: + fsdp_plugin.set_mixed_precision(self._mixed_precision) + self.fsdp_plugin = fsdp_plugin + if os.environ.get( + "ACCELERATE_USE_MEGATRON_LM", "false" + ).lower() == "true" and self.distributed_type not in [ + DistributedType.MULTI_XPU, + ]: + self.distributed_type = DistributedType.MEGATRON_LM + megatron_lm_plugin.set_mixed_precision(self._mixed_precision) + self.megatron_lm_plugin = megatron_lm_plugin + elif self.distributed_type in [DistributedType.MULTI_CPU, DistributedType.MULTI_XPU, DistributedType.NO]: + if is_ipex_available(): + # check if user disables it explicitly + self.use_ipex = parse_flag_from_env("ACCELERATE_USE_IPEX", default=True) + else: + self.use_ipex = False + if ( + self.dynamo_plugin.backend != DynamoBackend.NO + and self._mixed_precision == "no" + and self.device.type == "cuda" + ): + torch.backends.cuda.matmul.allow_tf32 = True + if ( + self.dynamo_plugin.backend != DynamoBackend.NO + and self._mixed_precision == "no" + and self.device.type == "musa" + ): + torch.backends.musa.matmul.allow_tf32 = True + PartialState._shared_state["distributed_type"] = self.distributed_type + + @property + def initialized(self) -> bool: + return self._shared_state != PartialState._shared_state + + def __repr__(self): + repr = PartialState().__repr__() + f"\nMixed precision type: {self.mixed_precision}\n" + if self.distributed_type == DistributedType.DEEPSPEED: + repr += f"ds_config: {self.deepspeed_plugin.deepspeed_config}\n" + return repr + + def _check_initialized(self, mixed_precision=None, cpu=None): + "Checks if a modification is trying to be made and the `AcceleratorState` has already been initialized" + if self.initialized: + err = "AcceleratorState has already been initialized and cannot be changed, restart your runtime completely and pass `{flag}` to `Accelerator()`." + if cpu and self.device.type != "cpu": + raise ValueError(err.format(flag="cpu=True")) + if ( + mixed_precision is not None + and mixed_precision != self._mixed_precision + and self.distributed_type != DistributedType.DEEPSPEED + ): + raise ValueError(err.format(flag=f"mixed_precision='{mixed_precision}'")) + + @property + def mixed_precision(self): + if self.distributed_type == DistributedType.DEEPSPEED and self._mixed_precision != "fp8": + config = self.deepspeed_plugin.deepspeed_config + if config.get("fp16", {}).get("enabled", False): + mixed_precision = "fp16" + elif config.get("bf16", {}).get("enabled", False): + mixed_precision = "bf16" + else: + mixed_precision = "no" + else: + mixed_precision = self._mixed_precision + return mixed_precision + + @staticmethod + def _reset_state(reset_partial_state: bool = False): + "Resets `_shared_state`, is used internally and should not be called" + AcceleratorState._shared_state.clear() + if reset_partial_state: + PartialState._reset_state() + + def destroy_process_group(self, group=None): + """ + Destroys the process group. If one is not specified, the default process group is destroyed. + + If `self.fork_launched` is `True` and `group` is `None`, nothing happens. + """ + PartialState().destroy_process_group(group) + + @property + def fork_launched(self): + return PartialState().fork_launched + + @property + def use_distributed(self): + """ + Whether the Accelerator is configured for distributed training + """ + return PartialState().use_distributed + + @property + def is_fsdp2(self) -> bool: + return self.distributed_type == DistributedType.FSDP and self.fsdp_plugin.fsdp_version == 2 + + @property + def is_last_process(self) -> bool: + "Returns whether the current process is the last one" + return PartialState().is_last_process + + @property + def is_main_process(self) -> bool: + "Returns whether the current process is the main process" + return PartialState().is_main_process + + @property + def is_local_main_process(self) -> bool: + "Returns whether the current process is the main process on the local node" + return PartialState().is_local_main_process + + def wait_for_everyone(self): + PartialState().wait_for_everyone() + + @contextmanager + def split_between_processes(self, inputs: list | tuple | dict | torch.Tensor, apply_padding: bool = False): + """ + Splits `input` between `self.num_processes` quickly and can be then used on that process. Useful when doing + distributed inference, such as with different prompts. + + Note that when using a `dict`, all keys need to have the same number of elements. + + Args: + inputs (`list`, `tuple`, `torch.Tensor`, or `dict` of `list`/`tuple`/`torch.Tensor`): + The input to split between processes. + apply_padding (`bool`, `optional`, defaults to `False`): + Whether to apply padding by repeating the last element of the input so that all processes have the same + number of elements. Useful when trying to perform actions such as `gather()` on the outputs or passing + in less inputs than there are processes. If so, just remember to drop the padded elements afterwards. + + + Example: + + ```python + # Assume there are two processes + from accelerate.state import AcceleratorState + + state = AcceleratorState() + with state.split_between_processes(["A", "B", "C"]) as inputs: + print(inputs) + # Process 0 + ["A", "B"] + # Process 1 + ["C"] + + with state.split_between_processes(["A", "B", "C"], apply_padding=True) as inputs: + print(inputs) + # Process 0 + ["A", "B"] + # Process 1 + ["C", "C"] + ``` + """ + with PartialState().split_between_processes(inputs, apply_padding=apply_padding) as inputs: + yield inputs + + @contextmanager + def main_process_first(self): + """ + Lets the main process go first inside a with block. + + The other processes will enter the with block after the main process exits. + """ + with PartialState().main_process_first(): + yield + + @contextmanager + def local_main_process_first(self): + """ + Lets the local main process go inside a with block. + + The other processes will enter the with block after the main process exits. + """ + with PartialState().local_main_process_first(): + yield + + @property + def deepspeed_plugin(self): + """ + Returns the currently active DeepSpeedPlugin. + + If not using deepspeed, returns `None`. + """ + # To maintain original behavior, return None if not using deepspeed. + if self.distributed_type != DistributedType.DEEPSPEED: + return None + from accelerate.utils.deepspeed import get_active_deepspeed_plugin + + return get_active_deepspeed_plugin(self) + + @deepspeed_required + def get_deepspeed_plugin(self, name: str): + """ + Returns the DeepSpeedPlugin with the given plugin_key. + """ + return self.deepspeed_plugins[name] + + @deepspeed_required + def select_deepspeed_plugin(self, name: str | None = None): + """ + Activates the DeepSpeedPlugin with the given `name`, and will disable all other plugins. + """ + for key, plugin in self.deepspeed_plugins.items(): + if key != name: + plugin._unselect() + self.deepspeed_plugins[name].select(_from_accelerator_state=True) + + def print(self, *args, **kwargs): + PartialState().print(*args, **kwargs) + + def __getattr__(self, name: str): + # By this point we know that no attributes of `self` contain `name`, + # so we just modify the error message + if name in self._known_attrs: + raise AttributeError( + f"`AcceleratorState` object has no attribute `{name}`. " + "This happens if `AcceleratorState._reset_state()` was called and " + "an `Accelerator` or `PartialState` was not reinitialized." + ) + # Raise a typical AttributeError + raise AttributeError(f"'AcceleratorState' object has no attribute '{name}'") + + +class GradientState: + """ + Singleton class that has information related to gradient synchronization for gradient accumulation + + **Available attributes:** + + - **end_of_dataloader** (`bool`) -- Whether we have reached the end the current dataloader + - **remainder** (`int`) -- The number of extra samples that were added from padding the dataloader + - **sync_gradients** (`bool`) -- Whether the gradients should be synced across all devices + - **active_dataloader** (`Optional[DataLoader]`) -- The dataloader that is currently being iterated over + - **dataloader_references** (`List[Optional[DataLoader]]`) -- A list of references to the dataloaders that are + being iterated over + - **num_steps** (`int`) -- The number of steps to accumulate over + - **adjust_scheduler** (`bool`) -- Whether the scheduler should be adjusted to account for the gradient + accumulation + - **sync_with_dataloader** (`bool`) -- Whether the gradients should be synced at the end of the dataloader + iteration and the number of total steps reset + - **is_xla_gradients_synced** (`bool`) -- Whether the XLA gradients have been synchronized. It is initialized + as false. Once gradients have been reduced before the optimizer step, this flag is set to true. Subsequently, + after each step, the flag is reset to false. FSDP will always synchronize the gradients, hence + is_xla_gradients_synced is always true. + """ + + _shared_state = SharedDict() + + def __init__(self, gradient_accumulation_plugin: GradientAccumulationPlugin | None = None): + self.__dict__ = self._shared_state + if not self.initialized: + self.sync_gradients = True + self._dataloader_references_ref = [None] + self.plugin_kwargs = ( + gradient_accumulation_plugin.to_kwargs() if gradient_accumulation_plugin is not None else {} + ) + self._is_xla_gradients_synced = False + + # Plugin args are different and can be updated + if gradient_accumulation_plugin is not None and self.plugin_kwargs != gradient_accumulation_plugin.to_kwargs(): + self.plugin_kwargs = gradient_accumulation_plugin.to_kwargs() + + @property + def num_steps(self) -> int: + "Returns the number of steps to accumulate over" + return self.plugin_kwargs.get("num_steps", 1) + + @property + def adjust_scheduler(self) -> bool: + "Returns whether the scheduler should be adjusted" + return self.plugin_kwargs.get("adjust_scheduler", False) + + @property + def sync_with_dataloader(self) -> bool: + "Returns whether the gradients should be synced at the end of the dataloader iteration and the number of total steps reset" + return self.plugin_kwargs.get("sync_with_dataloader", True) + + @property + def initialized(self) -> bool: + "Returns whether the `GradientState` has been initialized" + return GradientState._shared_state != {} + + @property + def end_of_dataloader(self) -> bool: + "Returns whether we have reached the end of the current dataloader" + if not self.in_dataloader: + return False + return self.active_dataloader.end_of_dataloader + + @property + def remainder(self) -> int: + "Returns the number of extra samples that were added from padding the dataloader" + if not self.in_dataloader: + return -1 + return self.active_dataloader.remainder + + def __repr__(self): + return ( + f"Sync Gradients: {self.sync_gradients}\n" + f"At end of current dataloader: {self.end_of_dataloader}\n" + f"Extra samples added: {self.remainder}\n" + f"Gradient accumulation plugin: {self.plugin_kwargs}\n" + ) + + @property + def is_xla_gradients_synced(self): + "Returns the value of is_xla_gradients_synced. FSDP will always synchronize the gradients, hence is_xla_gradients_synced is always true." + if parse_flag_from_env("ACCELERATE_USE_FSDP", default=False): + return True + return self._is_xla_gradients_synced + + @is_xla_gradients_synced.setter + def is_xla_gradients_synced(self, is_synced): + "Set the _is_xla_gradients_synced attribute." + self._is_xla_gradients_synced = is_synced + + def _set_sync_gradients(self, sync_gradients): + "Private function that sets whether gradients should be synchronized. Users should not have to call this." + self.sync_gradients = sync_gradients + # Allow grad-sync to automatically work on TPUs + if ( + self.sync_gradients + and is_torch_xla_available(check_is_tpu=True) + and PartialState().distributed_type == DistributedType.XLA + ): + xm.mark_step() + + def _add_dataloader(self, dataloader): + "Private function that adds a dataloader to `self.dataloader_references` and sets `in_dataloader` to `True`. Users should not have to call this." + # We explicitly use assignment to ensure that the property setter is triggered, which is required for garbage collection. + # Avoid using self.dataloader_references.append as it will not trigger the setter. + self.dataloader_references += [dataloader] + + def _remove_dataloader(self, dataloader): + "Private function that removes a dataloader from `self.dataloader_references` and sets `in_dataloader` to `False` if there are no more dataloaders. Users should not have to call this." + # We explicitly use assignment to ensure that the property setter is triggered. + self.dataloader_references = [ + dataloader_ref for dataloader_ref in self.dataloader_references if dataloader_ref != dataloader + ] + + @property + def active_dataloader(self): + return self.dataloader_references[-1] + + @property + def dataloader_references(self): + # We use a property getter and setter with weakrefs to avoid circular references that prevent garbage collection + return [reference() if reference is not None else reference for reference in self._dataloader_references_ref] + + @dataloader_references.setter + def dataloader_references(self, references): + self._dataloader_references_ref = [ + weakref.ref(dataloader) if dataloader is not None else dataloader for dataloader in references + ] + + @property + def in_dataloader(self) -> bool: + "Returns whether the current process is in a dataloader" + return self.active_dataloader is not None + + @staticmethod + def _reset_state(): + "Resets `_shared_state`, is used internally and should not be called" + GradientState._shared_state.clear() diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/tracking.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/tracking.py new file mode 100644 index 0000000000000000000000000000000000000000..ecd722569fcde9956355dc6c81b96c1aac08ea20 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/tracking.py @@ -0,0 +1,1326 @@ +# Copyright 2022 The HuggingFace Team. All rights reserved. +# +# 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. + +# Expectation: +# Provide a project dir name, then each type of logger gets stored in project/{`logging_dir`} + +import json +import os +import time +from functools import wraps +from typing import Any, Optional, Union + +import yaml +from packaging import version + +from .logging import get_logger +from .state import PartialState +from .utils import ( + LoggerType, + compare_versions, + is_aim_available, + is_clearml_available, + is_comet_ml_available, + is_dvclive_available, + is_mlflow_available, + is_swanlab_available, + is_tensorboard_available, + is_trackio_available, + is_wandb_available, + listify, +) + + +_available_trackers = [] + +if is_tensorboard_available(): + _available_trackers.append(LoggerType.TENSORBOARD) + +if is_wandb_available(): + _available_trackers.append(LoggerType.WANDB) + +if is_comet_ml_available(): + _available_trackers.append(LoggerType.COMETML) + +if is_aim_available(): + _available_trackers.append(LoggerType.AIM) + +if is_mlflow_available(): + _available_trackers.append(LoggerType.MLFLOW) + +if is_clearml_available(): + _available_trackers.append(LoggerType.CLEARML) + +if is_dvclive_available(): + _available_trackers.append(LoggerType.DVCLIVE) + +if is_swanlab_available(): + _available_trackers.append(LoggerType.SWANLAB) + +if is_trackio_available(): + _available_trackers.append(LoggerType.TRACKIO) + +logger = get_logger(__name__) + + +def on_main_process(function): + """ + Decorator to selectively run the decorated function on the main process only based on the `main_process_only` + attribute in a class. + + Checks at function execution rather than initialization time, not triggering the initialization of the + `PartialState`. + """ + + @wraps(function) + def execute_on_main_process(self, *args, **kwargs): + if getattr(self, "main_process_only", False): + return PartialState().on_main_process(function)(self, *args, **kwargs) + else: + return function(self, *args, **kwargs) + + return execute_on_main_process + + +def get_available_trackers(): + "Returns a list of all supported available trackers in the system" + return _available_trackers + + +class GeneralTracker: + """ + A base Tracker class to be used for all logging integration implementations. + + Each function should take in `**kwargs` that will automatically be passed in from a base dictionary provided to + [`Accelerator`]. + + Should implement `name`, `requires_logging_directory`, and `tracker` properties such that: + + `name` (`str`): String representation of the tracker class name, such as "TensorBoard" `requires_logging_directory` + (`bool`): Whether the logger requires a directory to store their logs. `tracker` (`object`): Should return internal + tracking mechanism used by a tracker class (such as the `run` for wandb) + + Implementations can also include a `main_process_only` (`bool`) attribute to toggle if relevant logging, init, and + other functions should occur on the main process or across all processes (by default will use `True`) + """ + + main_process_only = True + + def __init__(self, _blank=False): + if not _blank: + err = "" + if not hasattr(self, "name"): + err += "`name`" + if not hasattr(self, "requires_logging_directory"): + if len(err) > 0: + err += ", " + err += "`requires_logging_directory`" + + # as tracker is a @property that relies on post-init + if "tracker" not in dir(self): + if len(err) > 0: + err += ", " + err += "`tracker`" + if len(err) > 0: + raise NotImplementedError( + f"The implementation for this tracker class is missing the following " + f"required attributes. Please define them in the class definition: " + f"{err}" + ) + + def start(self): + """ + Lazy initialization of the tracker inside Accelerator to avoid initializing PartialState before + InitProcessGroupKwargs. + """ + pass + + def store_init_configuration(self, values: dict): + """ + Logs `values` as hyperparameters for the run. Implementations should use the experiment configuration + functionality of a tracking API. + + Args: + values (Dictionary `str` to `bool`, `str`, `float` or `int`): + Values to be stored as initial hyperparameters as key-value pairs. The values need to have type `bool`, + `str`, `float`, `int`, or `None`. + """ + pass + + def log(self, values: dict, step: Optional[int], **kwargs): + """ + Logs `values` to the current run. Base `log` implementations of a tracking API should go in here, along with + special behavior for the `step parameter. + + Args: + values (Dictionary `str` to `str`, `float`, or `int`): + Values to be logged as key-value pairs. The values need to have type `str`, `float`, or `int`. + step (`int`, *optional*): + The run step. If included, the log will be affiliated with this step. + """ + pass + + def finish(self): + """ + Should run any finalizing functions within the tracking API. If the API should not have one, just don't + overwrite that method. + """ + pass + + +class TensorBoardTracker(GeneralTracker): + """ + A `Tracker` class that supports `tensorboard`. Should be initialized at the start of your script. + + Args: + run_name (`str`): + The name of the experiment run + logging_dir (`str`, `os.PathLike`): + Location for TensorBoard logs to be stored. + **kwargs (additional keyword arguments, *optional*): + Additional key word arguments passed along to the `tensorboard.SummaryWriter.__init__` method. + """ + + name = "tensorboard" + requires_logging_directory = True + + def __init__(self, run_name: str, logging_dir: Union[str, os.PathLike], **kwargs): + super().__init__() + self.run_name = run_name + self.logging_dir_param = logging_dir + self.init_kwargs = kwargs + + @on_main_process + def start(self): + try: + from torch.utils import tensorboard + except ModuleNotFoundError: + import tensorboardX as tensorboard + self.logging_dir = os.path.join(self.logging_dir_param, self.run_name) + self.writer = tensorboard.SummaryWriter(self.logging_dir, **self.init_kwargs) + logger.debug(f"Initialized TensorBoard project {self.run_name} logging to {self.logging_dir}") + logger.debug( + "Make sure to log any initial configurations with `self.store_init_configuration` before training!" + ) + + @property + def tracker(self): + return self.writer + + @on_main_process + def store_init_configuration(self, values: dict): + """ + Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment. Stores the + hyperparameters in a yaml file for future use. + + Args: + values (Dictionary `str` to `bool`, `str`, `float` or `int`): + Values to be stored as initial hyperparameters as key-value pairs. The values need to have type `bool`, + `str`, `float`, `int`, or `None`. + """ + self.writer.add_hparams(values, metric_dict={}) + self.writer.flush() + project_run_name = time.time() + dir_name = os.path.join(self.logging_dir, str(project_run_name)) + os.makedirs(dir_name, exist_ok=True) + with open(os.path.join(dir_name, "hparams.yml"), "w") as outfile: + try: + yaml.dump(values, outfile) + except yaml.representer.RepresenterError: + logger.error("Serialization to store hyperparameters failed") + raise + logger.debug("Stored initial configuration hyperparameters to TensorBoard and hparams yaml file") + + @on_main_process + def log(self, values: dict, step: Optional[int] = None, **kwargs): + """ + Logs `values` to the current run. + + Args: + values (Dictionary `str` to `str`, `float`, `int` or `dict` of `str` to `float`/`int`): + Values to be logged as key-value pairs. The values need to have type `str`, `float`, `int` or `dict` of + `str` to `float`/`int`. + step (`int`, *optional*): + The run step. If included, the log will be affiliated with this step. + kwargs: + Additional key word arguments passed along to either `SummaryWriter.add_scaler`, + `SummaryWriter.add_text`, or `SummaryWriter.add_scalers` method based on the contents of `values`. + """ + values = listify(values) + for k, v in values.items(): + if isinstance(v, (int, float)): + self.writer.add_scalar(k, v, global_step=step, **kwargs) + elif isinstance(v, str): + self.writer.add_text(k, v, global_step=step, **kwargs) + elif isinstance(v, dict): + self.writer.add_scalars(k, v, global_step=step, **kwargs) + self.writer.flush() + logger.debug("Successfully logged to TensorBoard") + + @on_main_process + def log_images(self, values: dict, step: Optional[int], **kwargs): + """ + Logs `images` to the current run. + + Args: + values (Dictionary `str` to `List` of `np.ndarray` or `PIL.Image`): + Values to be logged as key-value pairs. The values need to have type `List` of `np.ndarray` or + step (`int`, *optional*): + The run step. If included, the log will be affiliated with this step. + kwargs: + Additional key word arguments passed along to the `SummaryWriter.add_image` method. + """ + for k, v in values.items(): + self.writer.add_images(k, v, global_step=step, **kwargs) + logger.debug("Successfully logged images to TensorBoard") + + @on_main_process + def finish(self): + """ + Closes `TensorBoard` writer + """ + self.writer.close() + logger.debug("TensorBoard writer closed") + + +class WandBTracker(GeneralTracker): + """ + A `Tracker` class that supports `wandb`. Should be initialized at the start of your script. + + Args: + run_name (`str`): + The name of the experiment run. + **kwargs (additional keyword arguments, *optional*): + Additional key word arguments passed along to the `wandb.init` method. + """ + + name = "wandb" + requires_logging_directory = False + main_process_only = False + + def __init__(self, run_name: str, **kwargs): + super().__init__() + self.run_name = run_name + self.init_kwargs = kwargs + + @on_main_process + def start(self): + import wandb + + self.run = wandb.init(project=self.run_name, **self.init_kwargs) + logger.debug(f"Initialized WandB project {self.run_name}") + logger.debug( + "Make sure to log any initial configurations with `self.store_init_configuration` before training!" + ) + + @property + def tracker(self): + return self.run + + @on_main_process + def store_init_configuration(self, values: dict): + """ + Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment. + + Args: + values (Dictionary `str` to `bool`, `str`, `float` or `int`): + Values to be stored as initial hyperparameters as key-value pairs. The values need to have type `bool`, + `str`, `float`, `int`, or `None`. + """ + import wandb + + if os.environ.get("WANDB_MODE") == "offline": + # In offline mode, restart wandb with config included + if hasattr(self, "run") and self.run: + self.run.finish() + + init_kwargs = self.init_kwargs.copy() + init_kwargs["config"] = values + self.run = wandb.init(project=self.run_name, **init_kwargs) + else: + wandb.config.update(values, allow_val_change=True) + logger.debug("Stored initial configuration hyperparameters to WandB") + + @on_main_process + def log(self, values: dict, step: Optional[int] = None, **kwargs): + """ + Logs `values` to the current run. + + Args: + values (Dictionary `str` to `str`, `float`, `int` or `dict` of `str` to `float`/`int`): + Values to be logged as key-value pairs. The values need to have type `str`, `float`, `int` or `dict` of + `str` to `float`/`int`. + step (`int`, *optional*): + The run step. If included, the log will be affiliated with this step. + kwargs: + Additional key word arguments passed along to the `wandb.log` method. + """ + self.run.log(values, step=step, **kwargs) + logger.debug("Successfully logged to WandB") + + @on_main_process + def log_images(self, values: dict, step: Optional[int] = None, **kwargs): + """ + Logs `images` to the current run. + + Args: + values (Dictionary `str` to `List` of `np.ndarray` or `PIL.Image`): + Values to be logged as key-value pairs. The values need to have type `List` of `np.ndarray` or + step (`int`, *optional*): + The run step. If included, the log will be affiliated with this step. + kwargs: + Additional key word arguments passed along to the `wandb.log` method. + """ + import wandb + + for k, v in values.items(): + self.log({k: [wandb.Image(image) for image in v]}, step=step, **kwargs) + logger.debug("Successfully logged images to WandB") + + @on_main_process + def log_table( + self, + table_name: str, + columns: Optional[list[str]] = None, + data: Optional[list[list[Any]]] = None, + dataframe: Any = None, + step: Optional[int] = None, + **kwargs, + ): + """ + Log a Table containing any object type (text, image, audio, video, molecule, html, etc). Can be defined either + with `columns` and `data` or with `dataframe`. + + Args: + table_name (`str`): + The name to give to the logged table on the wandb workspace + columns (list of `str`, *optional*): + The name of the columns on the table + data (List of List of Any data type, *optional*): + The data to be logged in the table + dataframe (Any data type, *optional*): + The data to be logged in the table + step (`int`, *optional*): + The run step. If included, the log will be affiliated with this step. + """ + import wandb + + values = {table_name: wandb.Table(columns=columns, data=data, dataframe=dataframe)} + self.log(values, step=step, **kwargs) + + @on_main_process + def finish(self): + """ + Closes `wandb` writer + """ + self.run.finish() + logger.debug("WandB run closed") + + +class TrackioTracker(GeneralTracker): + """ + A `Tracker` class that supports `trackio`. Should be initialized at the start of your script. + + Args: + run_name (`str`): + The name of the experiment run. Will be used as the `project` name when instantiating trackio. + **kwargs (additional keyword arguments, *optional*): + Additional key word arguments passed along to the `trackio.init` method. Refer to this + [init](https://github.com/gradio-app/trackio/blob/814809552310468b13f84f33764f1369b4e5136c/trackio/__init__.py#L22) + to see all supported key word arguments. + """ + + name = "trackio" + requires_logging_directory = False + main_process_only = False + + def __init__(self, run_name: str, **kwargs): + super().__init__() + self.run_name = run_name + self.init_kwargs = kwargs + + @on_main_process + def start(self): + import trackio + + self.run = trackio.init(project=self.run_name, **self.init_kwargs) + logger.debug(f"Initialized trackio project {self.run_name}") + logger.debug( + "Make sure to log any initial configurations with `self.store_init_configuration` before training!" + ) + + @property + def tracker(self): + return self.run + + @on_main_process + def store_init_configuration(self, values: dict): + """ + Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment. + + Args: + values (Dictionary `str` to `bool`, `str`, `float` or `int`): + Values to be stored as initial hyperparameters as key-value pairs. The values need to have type `bool`, + `str`, `float`, `int`, or `None`. + """ + import trackio + + trackio.config.update(values, allow_val_change=True) + logger.debug("Stored initial configuration hyperparameters to trackio") + + @on_main_process + def log(self, values: dict, step: Optional[int] = None, **kwargs): + """ + Logs `values` to the current run. + + Args: + values (Dictionary `str` to `str`, `float`, `int` or `dict` of `str` to `float`/`int`): + Values to be logged as key-value pairs. The values need to have type `str`, `float`, `int` or `dict` of + `str` to `float`/`int`. + step (`int`, *optional*): + The run step. If included, the log will be affiliated with this step. + kwargs: + Additional key word arguments passed along to the `trackio.log` method. + """ + self.run.log(values, **kwargs) + logger.debug("Successfully logged to trackio") + + @on_main_process + def finish(self): + """ + Closes `trackio` run + """ + self.run.finish() + logger.debug("trackio run closed") + + +class CometMLTracker(GeneralTracker): + """ + A `Tracker` class that supports `comet_ml`. Should be initialized at the start of your script. + + API keys must be stored in a Comet config file. + + Note: + For `comet_ml` versions < 3.41.0, additional keyword arguments are passed to `comet_ml.Experiment` instead: + https://www.comet.com/docs/v2/api-and-sdk/python-sdk/reference/Experiment/#comet_ml.Experiment.__init__ + + Args: + run_name (`str`): + The name of the experiment run. + **kwargs (additional keyword arguments, *optional*): + Additional key word arguments passed along to the `comet_ml.start` method: + https://www.comet.com/docs/v2/api-and-sdk/python-sdk/reference/start/ + """ + + name = "comet_ml" + requires_logging_directory = False + + def __init__(self, run_name: str, **kwargs): + super().__init__() + self.run_name = run_name + self.init_kwargs = kwargs + + @on_main_process + def start(self): + import comet_ml + + comet_version = version.parse(comet_ml.__version__) + if compare_versions(comet_version, ">=", "3.41.0"): + self.writer = comet_ml.start(project_name=self.run_name, **self.init_kwargs) + else: + logger.info("Update `comet_ml` (>=3.41.0) for experiment reuse and offline support.") + self.writer = comet_ml.Experiment(project_name=self.run_name, **self.init_kwargs) + + logger.debug(f"Initialized CometML project {self.run_name}") + logger.debug( + "Make sure to log any initial configurations with `self.store_init_configuration` before training!" + ) + + @property + def tracker(self): + return self.writer + + @on_main_process + def store_init_configuration(self, values: dict): + """ + Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment. + + Args: + values (Dictionary `str` to `bool`, `str`, `float` or `int`): + Values to be stored as initial hyperparameters as key-value pairs. The values need to have type `bool`, + `str`, `float`, `int`, or `None`. + """ + self.writer.log_parameters(values) + logger.debug("Stored initial configuration hyperparameters to Comet") + + @on_main_process + def log(self, values: dict, step: Optional[int] = None, **kwargs): + """ + Logs `values` to the current run. + + Args: + values (Dictionary `str` to `str`, `float`, `int` or `dict` of `str` to `float`/`int`): + Values to be logged as key-value pairs. The values need to have type `str`, `float`, `int` or `dict` of + `str` to `float`/`int`. + step (`int`, *optional*): + The run step. If included, the log will be affiliated with this step. + kwargs: + Additional key word arguments passed along to either `Experiment.log_metric`, `Experiment.log_other`, + or `Experiment.log_metrics` method based on the contents of `values`. + """ + if step is not None: + self.writer.set_step(step) + for k, v in values.items(): + if isinstance(v, (int, float)): + self.writer.log_metric(k, v, step=step, **kwargs) + elif isinstance(v, str): + self.writer.log_other(k, v, **kwargs) + elif isinstance(v, dict): + self.writer.log_metrics(v, step=step, **kwargs) + logger.debug("Successfully logged to Comet") + + @on_main_process + def finish(self): + """ + Flush `comet-ml` writer + """ + self.writer.end() + logger.debug("Comet run flushed") + + +class AimTracker(GeneralTracker): + """ + A `Tracker` class that supports `aim`. Should be initialized at the start of your script. + + Args: + run_name (`str`): + The name of the experiment run. + **kwargs (additional keyword arguments, *optional*): + Additional key word arguments passed along to the `Run.__init__` method. + """ + + name = "aim" + requires_logging_directory = True + + def __init__(self, run_name: str, logging_dir: Optional[Union[str, os.PathLike]] = ".", **kwargs): + super().__init__() + self.run_name = run_name + self.aim_repo_path = logging_dir + self.init_kwargs = kwargs + + @on_main_process + def start(self): + from aim import Run + + self.writer = Run(repo=self.aim_repo_path, **self.init_kwargs) + self.writer.name = self.run_name + logger.debug(f"Initialized Aim project {self.run_name}") + logger.debug( + "Make sure to log any initial configurations with `self.store_init_configuration` before training!" + ) + + @property + def tracker(self): + return self.writer + + @on_main_process + def store_init_configuration(self, values: dict): + """ + Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment. + + Args: + values (`dict`): + Values to be stored as initial hyperparameters as key-value pairs. + """ + self.writer["hparams"] = values + + @on_main_process + def log(self, values: dict, step: Optional[int], **kwargs): + """ + Logs `values` to the current run. + + Args: + values (`dict`): + Values to be logged as key-value pairs. + step (`int`, *optional*): + The run step. If included, the log will be affiliated with this step. + kwargs: + Additional key word arguments passed along to the `Run.track` method. + """ + # Note: replace this with the dictionary support when merged + for key, value in values.items(): + self.writer.track(value, name=key, step=step, **kwargs) + + @on_main_process + def log_images(self, values: dict, step: Optional[int] = None, kwargs: Optional[dict[str, dict]] = None): + """ + Logs `images` to the current run. + + Args: + values (`Dict[str, Union[np.ndarray, PIL.Image, Tuple[np.ndarray, str], Tuple[PIL.Image, str]]]`): + Values to be logged as key-value pairs. The values need to have type `np.ndarray` or PIL.Image. If a + tuple is provided, the first element should be the image and the second element should be the caption. + step (`int`, *optional*): + The run step. If included, the log will be affiliated with this step. + kwargs (`Dict[str, dict]`): + Additional key word arguments passed along to the `Run.Image` and `Run.track` method specified by the + keys `aim_image` and `track`, respectively. + """ + import aim + + aim_image_kw = {} + track_kw = {} + + if kwargs is not None: + aim_image_kw = kwargs.get("aim_image", {}) + track_kw = kwargs.get("track", {}) + + for key, value in values.items(): + if isinstance(value, tuple): + img, caption = value + else: + img, caption = value, "" + aim_image = aim.Image(img, caption=caption, **aim_image_kw) + self.writer.track(aim_image, name=key, step=step, **track_kw) + + @on_main_process + def finish(self): + """ + Closes `aim` writer + """ + self.writer.close() + + +class MLflowTracker(GeneralTracker): + """ + A `Tracker` class that supports `mlflow`. Should be initialized at the start of your script. + + Args: + experiment_name (`str`, *optional*): + Name of the experiment. Environment variable MLFLOW_EXPERIMENT_NAME has priority over this argument. + logging_dir (`str` or `os.PathLike`, defaults to `"."`): + Location for mlflow logs to be stored. + run_id (`str`, *optional*): + If specified, get the run with the specified UUID and log parameters and metrics under that run. The run’s + end time is unset and its status is set to running, but the run’s other attributes (source_version, + source_type, etc.) are not changed. Environment variable MLFLOW_RUN_ID has priority over this argument. + tags (`Dict[str, str]`, *optional*): + An optional `dict` of `str` keys and values, or a `str` dump from a `dict`, to set as tags on the run. If a + run is being resumed, these tags are set on the resumed run. If a new run is being created, these tags are + set on the new run. Environment variable MLFLOW_TAGS has priority over this argument. + nested_run (`bool`, *optional*, defaults to `False`): + Controls whether run is nested in parent run. True creates a nested run. Environment variable + MLFLOW_NESTED_RUN has priority over this argument. + run_name (`str`, *optional*): + Name of new run (stored as a mlflow.runName tag). Used only when `run_id` is unspecified. + description (`str`, *optional*): + An optional string that populates the description box of the run. If a run is being resumed, the + description is set on the resumed run. If a new run is being created, the description is set on the new + run. + """ + + name = "mlflow" + requires_logging_directory = False + + def __init__( + self, + experiment_name: Optional[str] = None, + logging_dir: Optional[Union[str, os.PathLike]] = None, + run_id: Optional[str] = None, + tags: Optional[Union[dict[str, Any], str]] = None, + nested_run: Optional[bool] = False, + run_name: Optional[str] = None, + description: Optional[str] = None, + ): + experiment_name = os.environ.get("MLFLOW_EXPERIMENT_NAME", experiment_name) + run_id = os.environ.get("MLFLOW_RUN_ID", run_id) + tags = os.environ.get("MLFLOW_TAGS", tags) + if isinstance(tags, str): + tags = json.loads(tags) + + nested_run = os.environ.get("MLFLOW_NESTED_RUN", nested_run) + + self.experiment_name = experiment_name + self.logging_dir = logging_dir + self.run_id = run_id + self.tags = tags + self.nested_run = nested_run + self.run_name = run_name + self.description = description + + @on_main_process + def start(self): + import mlflow + + exps = mlflow.search_experiments(filter_string=f"name = '{self.experiment_name}'") + if len(exps) > 0: + if len(exps) > 1: + logger.warning("Multiple experiments with the same name found. Using first one.") + experiment_id = exps[0].experiment_id + else: + experiment_id = mlflow.create_experiment( + name=self.experiment_name, + artifact_location=self.logging_dir, + tags=self.tags, + ) + + self.active_run = mlflow.start_run( + run_id=self.run_id, + experiment_id=experiment_id, + run_name=self.run_name, + nested=self.nested_run, + tags=self.tags, + description=self.description, + ) + + logger.debug(f"Initialized mlflow experiment {self.experiment_name}") + logger.debug( + "Make sure to log any initial configurations with `self.store_init_configuration` before training!" + ) + + @property + def tracker(self): + return self.active_run + + @on_main_process + def store_init_configuration(self, values: dict): + """ + Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment. + + Args: + values (`dict`): + Values to be stored as initial hyperparameters as key-value pairs. + """ + import mlflow + + for name, value in list(values.items()): + # internally, all values are converted to str in MLflow + if len(str(value)) > mlflow.utils.validation.MAX_PARAM_VAL_LENGTH: + logger.warning_once( + f'Accelerate is attempting to log a value of "{value}" for key "{name}" as a parameter. MLflow\'s' + f" log_param() only accepts values no longer than {mlflow.utils.validation.MAX_PARAM_VAL_LENGTH} characters so we dropped this attribute." + ) + del values[name] + + values_list = list(values.items()) + + # MLflow cannot log more than 100 values in one go, so we have to split it + for i in range(0, len(values_list), mlflow.utils.validation.MAX_PARAMS_TAGS_PER_BATCH): + mlflow.log_params(dict(values_list[i : i + mlflow.utils.validation.MAX_PARAMS_TAGS_PER_BATCH])) + + logger.debug("Stored initial configuration hyperparameters to MLflow") + + @on_main_process + def log(self, values: dict, step: Optional[int]): + """ + Logs `values` to the current run. + + Args: + values (`dict`): + Values to be logged as key-value pairs. + step (`int`, *optional*): + The run step. If included, the log will be affiliated with this step. + """ + metrics = {} + for k, v in values.items(): + if isinstance(v, (int, float)): + metrics[k] = v + else: + logger.warning_once( + f'MLflowTracker is attempting to log a value of "{v}" of type {type(v)} for key "{k}" as a metric. ' + "MLflow's log_metric() only accepts float and int types so we dropped this attribute." + ) + import mlflow + + mlflow.log_metrics(metrics, step=step) + logger.debug("Successfully logged to mlflow") + + @on_main_process + def log_figure(self, figure: Any, artifact_file: str, **save_kwargs): + """ + Logs an figure to the current run. + + Args: + figure (Any): + The figure to be logged. + artifact_file (`str`, *optional*): + The run-relative artifact file path in posixpath format to which the image is saved. + If not provided, the image is saved to a default location. + **kwargs: + Additional keyword arguments passed to the underlying mlflow.log_image function. + """ + import mlflow + + mlflow.log_figure(figure=figure, artifact_file=artifact_file, **save_kwargs) + logger.debug("Successfully logged image to mlflow") + + @on_main_process + def log_artifacts(self, local_dir: str, artifact_path: Optional[str] = None): + """ + Logs an artifacts (all content of a dir) to the current run. + + local_dir (`str`): + Path to the directory to be logged as an artifact. + artifact_path (`str`, *optional*): + Directory within the run's artifact directory where the artifact will be logged. If omitted, the + artifact will be logged to the root of the run's artifact directory. The run step. If included, the + artifact will be affiliated with this step. + """ + import mlflow + + mlflow.log_artifacts(local_dir=local_dir, artifact_path=artifact_path) + logger.debug("Successfully logged artofact to mlflow") + + @on_main_process + def log_artifact(self, local_path: str, artifact_path: Optional[str] = None): + """ + Logs an artifact (file) to the current run. + + local_path (`str`): + Path to the file to be logged as an artifact. + artifact_path (`str`, *optional*): + Directory within the run's artifact directory where the artifact will be logged. If omitted, the + artifact will be logged to the root of the run's artifact directory. The run step. If included, the + artifact will be affiliated with this step. + """ + import mlflow + + mlflow.log_artifact(local_path=local_path, artifact_path=artifact_path) + logger.debug("Successfully logged artofact to mlflow") + + @on_main_process + def finish(self): + """ + End the active MLflow run. + """ + import mlflow + + mlflow.end_run() + + +class ClearMLTracker(GeneralTracker): + """ + A `Tracker` class that supports `clearml`. Should be initialized at the start of your script. + + Args: + run_name (`str`, *optional*): + Name of the experiment. Environment variables `CLEARML_PROJECT` and `CLEARML_TASK` have priority over this + argument. + **kwargs (additional keyword arguments, *optional*): + Kwargs passed along to the `Task.__init__` method. + """ + + name = "clearml" + requires_logging_directory = False + + def __init__(self, run_name: Optional[str] = None, **kwargs): + super().__init__() + self.user_provided_run_name = run_name + self._initialized_externally = False + self.init_kwargs = kwargs + + @on_main_process + def start(self): + from clearml import Task + + current_task = Task.current_task() + if current_task: + self._initialized_externally = True + self.task = current_task + return + + task_init_args = {**self.init_kwargs} + task_init_args.setdefault("project_name", os.environ.get("CLEARML_PROJECT", self.user_provided_run_name)) + task_init_args.setdefault("task_name", os.environ.get("CLEARML_TASK", self.user_provided_run_name)) + self.task = Task.init(**task_init_args) + + @property + def tracker(self): + return self.task + + @on_main_process + def store_init_configuration(self, values: dict): + """ + Connect configuration dictionary to the Task object. Should be run at the beginning of your experiment. + + Args: + values (`dict`): + Values to be stored as initial hyperparameters as key-value pairs. + """ + return self.task.connect_configuration(values) + + @on_main_process + def log(self, values: dict[str, Union[int, float]], step: Optional[int] = None, **kwargs): + """ + Logs `values` dictionary to the current run. The dictionary keys must be strings. The dictionary values must be + ints or floats + + Args: + values (`Dict[str, Union[int, float]]`): + Values to be logged as key-value pairs. If the key starts with 'eval_'/'test_'/'train_', the value will + be reported under the 'eval'/'test'/'train' series and the respective prefix will be removed. + Otherwise, the value will be reported under the 'train' series, and no prefix will be removed. + step (`int`, *optional*): + If specified, the values will be reported as scalars, with the iteration number equal to `step`. + Otherwise they will be reported as single values. + kwargs: + Additional key word arguments passed along to the `clearml.Logger.report_single_value` or + `clearml.Logger.report_scalar` methods. + """ + clearml_logger = self.task.get_logger() + for k, v in values.items(): + if not isinstance(v, (int, float)): + logger.warning_once( + "Accelerator is attempting to log a value of " + f'"{v}" of type {type(v)} for key "{k}" as a scalar. ' + "This invocation of ClearML logger's report_scalar() " + "is incorrect so we dropped this attribute." + ) + continue + if step is None: + clearml_logger.report_single_value(name=k, value=v, **kwargs) + continue + title, series = ClearMLTracker._get_title_series(k) + clearml_logger.report_scalar(title=title, series=series, value=v, iteration=step, **kwargs) + + @on_main_process + def log_images(self, values: dict, step: Optional[int] = None, **kwargs): + """ + Logs `images` to the current run. + + Args: + values (`Dict[str, List[Union[np.ndarray, PIL.Image]]`): + Values to be logged as key-value pairs. The values need to have type `List` of `np.ndarray` or + step (`int`, *optional*): + The run step. If included, the log will be affiliated with this step. + kwargs: + Additional key word arguments passed along to the `clearml.Logger.report_image` method. + """ + clearml_logger = self.task.get_logger() + for k, v in values.items(): + title, series = ClearMLTracker._get_title_series(k) + clearml_logger.report_image(title=title, series=series, iteration=step, image=v, **kwargs) + + @on_main_process + def log_table( + self, + table_name: str, + columns: Optional[list[str]] = None, + data: Optional[list[list[Any]]] = None, + dataframe: Any = None, + step: Optional[int] = None, + **kwargs, + ): + """ + Log a Table to the task. Can be defined eitherwith `columns` and `data` or with `dataframe`. + + Args: + table_name (`str`): + The name of the table + columns (list of `str`, *optional*): + The name of the columns on the table + data (List of List of Any data type, *optional*): + The data to be logged in the table. If `columns` is not specified, then the first entry in data will be + the name of the columns of the table + dataframe (Any data type, *optional*): + The data to be logged in the table + step (`int`, *optional*): + The run step. If included, the log will be affiliated with this step. + kwargs: + Additional key word arguments passed along to the `clearml.Logger.report_table` method. + """ + to_report = dataframe + if dataframe is None: + if data is None: + raise ValueError( + "`ClearMLTracker.log_table` requires that `data` to be supplied if `dataframe` is `None`" + ) + to_report = [columns] + data if columns else data + title, series = ClearMLTracker._get_title_series(table_name) + self.task.get_logger().report_table(title=title, series=series, table_plot=to_report, iteration=step, **kwargs) + + @on_main_process + def finish(self): + """ + Close the ClearML task. If the task was initialized externally (e.g. by manually calling `Task.init`), this + function is a noop + """ + if self.task and not self._initialized_externally: + self.task.close() + + @staticmethod + def _get_title_series(name): + for prefix in ["eval", "test", "train"]: + if name.startswith(prefix + "_"): + return name[len(prefix) + 1 :], prefix + return name, "train" + + +class DVCLiveTracker(GeneralTracker): + """ + A `Tracker` class that supports `dvclive`. Should be initialized at the start of your script. + + Args: + run_name (`str`, *optional*): + Ignored for dvclive. See `kwargs` instead. + kwargs: + Additional key word arguments passed along to [`dvclive.Live()`](https://dvc.org/doc/dvclive/live). + + Example: + + ```py + from accelerate import Accelerator + + accelerator = Accelerator(log_with="dvclive") + accelerator.init_trackers(project_name="my_project", init_kwargs={"dvclive": {"dir": "my_directory"}}) + ``` + """ + + name = "dvclive" + requires_logging_directory = False + + def __init__(self, run_name: Optional[str] = None, live: Optional[Any] = None, **kwargs): + super().__init__() + self.live = live + self.init_kwargs = kwargs + + @on_main_process + def start(self): + from dvclive import Live + + self.live = self.live if self.live is not None else Live(**self.init_kwargs) + + @property + def tracker(self): + return self.live + + @on_main_process + def store_init_configuration(self, values: dict): + """ + Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment. Stores the + hyperparameters in a yaml file for future use. + + Args: + values (Dictionary `str` to `bool`, `str`, `float`, `int`, or a List or Dict of those types): + Values to be stored as initial hyperparameters as key-value pairs. The values need to have type `bool`, + `str`, `float`, or `int`. + """ + self.live.log_params(values) + + @on_main_process + def log(self, values: dict, step: Optional[int] = None, **kwargs): + """ + Logs `values` to the current run. + + Args: + values (Dictionary `str` to `str`, `float`, or `int`): + Values to be logged as key-value pairs. The values need to have type `str`, `float`, or `int`. + step (`int`, *optional*): + The run step. If included, the log will be affiliated with this step. + kwargs: + Additional key word arguments passed along to `dvclive.Live.log_metric()`. + """ + from dvclive.plots import Metric + + if step is not None: + self.live.step = step + for k, v in values.items(): + if Metric.could_log(v): + self.live.log_metric(k, v, **kwargs) + else: + logger.warning_once( + "Accelerator attempted to log a value of " + f'"{v}" of type {type(v)} for key "{k}" as a scalar. ' + "This invocation of DVCLive's Live.log_metric() " + "is incorrect so we dropped this attribute." + ) + self.live.next_step() + + @on_main_process + def finish(self): + """ + Closes `dvclive.Live()`. + """ + self.live.end() + + +class SwanLabTracker(GeneralTracker): + """ + A `Tracker` class that supports `swanlab`. Should be initialized at the start of your script. + + Args: + run_name (`str`): + The name of the experiment run. + **kwargs (additional keyword arguments, *optional*): + Additional key word arguments passed along to the `swanlab.init` method. + """ + + name = "swanlab" + requires_logging_directory = False + main_process_only = False + + def __init__(self, run_name: str, **kwargs): + super().__init__() + self.run_name = run_name + self.init_kwargs = kwargs + + @on_main_process + def start(self): + import swanlab + + self.run = swanlab.init(project=self.run_name, **self.init_kwargs) + swanlab.config["FRAMEWORK"] = "🤗Accelerate" # add accelerate logo in config + logger.debug(f"Initialized SwanLab project {self.run_name}") + logger.debug( + "Make sure to log any initial configurations with `self.store_init_configuration` before training!" + ) + + @property + def tracker(self): + return self.run + + @on_main_process + def store_init_configuration(self, values: dict): + """ + Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment. + + Args: + values (Dictionary `str` to `bool`, `str`, `float` or `int`): + Values to be stored as initial hyperparameters as key-value pairs. The values need to have type `bool`, + `str`, `float`, `int`, or `None`. + """ + import swanlab + + swanlab.config.update(values, allow_val_change=True) + logger.debug("Stored initial configuration hyperparameters to SwanLab") + + @on_main_process + def log(self, values: dict, step: Optional[int] = None, **kwargs): + """ + Logs `values` to the current run. + + Args: + data : Dict[str, DataType] + Data must be a dict. The key must be a string with 0-9, a-z, A-Z, " ", "_", "-", "/". The value must be a + `float`, `float convertible object`, `int` or `swanlab.data.BaseType`. + step : int, optional + The step number of the current data, if not provided, it will be automatically incremented. + If step is duplicated, the data will be ignored. + kwargs: + Additional key word arguments passed along to the `swanlab.log` method. Likes: + print_to_console : bool, optional + Whether to print the data to the console, the default is False. + """ + self.run.log(values, step=step, **kwargs) + logger.debug("Successfully logged to SwanLab") + + @on_main_process + def log_images(self, values: dict, step: Optional[int] = None, **kwargs): + """ + Logs `images` to the current run. + + Args: + values (Dictionary `str` to `List` of `np.ndarray` or `PIL.Image`): + Values to be logged as key-value pairs. The values need to have type `List` of `np.ndarray` or + step (`int`, *optional*): + The run step. If included, the log will be affiliated with this step. + kwargs: + Additional key word arguments passed along to the `swanlab.log` method. Likes: + print_to_console : bool, optional + Whether to print the data to the console, the default is False. + """ + import swanlab + + for k, v in values.items(): + self.log({k: [swanlab.Image(image) for image in v]}, step=step, **kwargs) + logger.debug("Successfully logged images to SwanLab") + + @on_main_process + def finish(self): + """ + Closes `swanlab` writer + """ + self.run.finish() + logger.debug("SwanLab run closed") + + +LOGGER_TYPE_TO_CLASS = { + "aim": AimTracker, + "comet_ml": CometMLTracker, + "mlflow": MLflowTracker, + "tensorboard": TensorBoardTracker, + "wandb": WandBTracker, + "clearml": ClearMLTracker, + "dvclive": DVCLiveTracker, + "swanlab": SwanLabTracker, + "trackio": TrackioTracker, +} + + +def filter_trackers( + log_with: list[Union[str, LoggerType, GeneralTracker]], + logging_dir: Optional[Union[str, os.PathLike]] = None, +): + """ + Takes in a list of potential tracker types and checks that: + - The tracker wanted is available in that environment + - Filters out repeats of tracker types + - If `all` is in `log_with`, will return all trackers in the environment + - If a tracker requires a `logging_dir`, ensures that `logging_dir` is not `None` + + Args: + log_with (list of `str`, [`~utils.LoggerType`] or [`~tracking.GeneralTracker`], *optional*): + A list of loggers to be setup for experiment tracking. Should be one or several of: + + - `"all"` + - `"tensorboard"` + - `"wandb"` + - `"trackio"` + - `"aim"` + - `"comet_ml"` + - `"mlflow"` + - `"dvclive"` + - `"swanlab"` + If `"all"` is selected, will pick up all available trackers in the environment and initialize them. Can + also accept implementations of `GeneralTracker` for custom trackers, and can be combined with `"all"`. + logging_dir (`str`, `os.PathLike`, *optional*): + A path to a directory for storing logs of locally-compatible loggers. + """ + loggers = [] + if log_with is not None: + if not isinstance(log_with, (list, tuple)): + log_with = [log_with] + if "all" in log_with or LoggerType.ALL in log_with: + loggers = [o for o in log_with if issubclass(type(o), GeneralTracker)] + get_available_trackers() + else: + for log_type in log_with: + if log_type not in LoggerType and not issubclass(type(log_type), GeneralTracker): + raise ValueError(f"Unsupported logging capability: {log_type}. Choose between {LoggerType.list()}") + if issubclass(type(log_type), GeneralTracker): + loggers.append(log_type) + else: + log_type = LoggerType(log_type) + if log_type not in loggers: + if log_type in get_available_trackers(): + tracker_init = LOGGER_TYPE_TO_CLASS[str(log_type)] + if tracker_init.requires_logging_directory: + if logging_dir is None: + raise ValueError( + f"Logging with `{log_type}` requires a `logging_dir` to be passed in." + ) + loggers.append(log_type) + else: + logger.debug(f"Tried adding logger {log_type}, but package is unavailable in the system.") + + return loggers diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/cuda_pathfinder-1.4.0.dist-info/INSTALLER b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/cuda_pathfinder-1.4.0.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/cuda_pathfinder-1.4.0.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/cuda_pathfinder-1.4.0.dist-info/METADATA b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/cuda_pathfinder-1.4.0.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..cb48899f6bb6bb8e75fec9f880beee1817522098 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/cuda_pathfinder-1.4.0.dist-info/METADATA @@ -0,0 +1,47 @@ +Metadata-Version: 2.4 +Name: cuda-pathfinder +Version: 1.4.0 +Summary: Pathfinder for CUDA components +Author-email: NVIDIA Corporation +License-Expression: Apache-2.0 +Project-URL: Repository, https://github.com/NVIDIA/cuda-python +Project-URL: Documentation, https://nvidia.github.io/cuda-python/ +Requires-Python: >=3.10 +Description-Content-Type: text/x-rst +License-File: LICENSE +Dynamic: license-file + +.. SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +.. SPDX-License-Identifier: Apache-2.0 + +******************************************************* +cuda-pathfinder: Utilities for locating CUDA components +******************************************************* + +.. image:: https://img.shields.io/badge/NVIDIA-black?logo=nvidia + :target: https://www.nvidia.com/ + :alt: NVIDIA + +`cuda.pathfinder `_ +aims to be a one-stop solution for locating CUDA components. Currently +it supports locating and loading dynamic libraries (``.so``, ``.dll``), and +locating CTK header directories. Support for other artifacts is in progress. + +* `Documentation `_ +* `Releases `_ +* `Repository `_ +* `Issue tracker `_ (select component ``cuda.pathfinder``) + +``cuda.pathfinder`` is under active development. Feedback and suggestions are welcome. + + +Installation +============ + +.. code-block:: bash + + pip install cuda-pathfinder + +``cuda-pathfinder`` is `CUDA Toolkit (CTK) `_ +version-agnostic. It follows the general CUDA Toolkit support policy: the +two most recent major versions are supported simultaneously. diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/cuda_pathfinder-1.4.0.dist-info/RECORD b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/cuda_pathfinder-1.4.0.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..15ae8d3579df7af2b76abf21621b6f261aa2c1aa --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/cuda_pathfinder-1.4.0.dist-info/RECORD @@ -0,0 +1,47 @@ +cuda/pathfinder/README.md,sha256=JPbUfq3xYFMGa1ZlnpHkKsGXjJi4Zr1nNU3SzWJIviw,123 +cuda/pathfinder/__init__.py,sha256=rHJIVMBSxsxzk46Z7cwZpjTyJ_fsQy7TOn4VlF8bgqs,3238 +cuda/pathfinder/__pycache__/__init__.cpython-312.pyc,, +cuda/pathfinder/__pycache__/_version.cpython-312.pyc,, +cuda/pathfinder/_binaries/__pycache__/find_nvidia_binary_utility.cpython-312.pyc,, +cuda/pathfinder/_binaries/__pycache__/supported_nvidia_binaries.cpython-312.pyc,, +cuda/pathfinder/_binaries/find_nvidia_binary_utility.py,sha256=8wrArz138JmogO9_ZA968OqQAf0rPkLYo9tndVUt4ug,4296 +cuda/pathfinder/_binaries/supported_nvidia_binaries.py,sha256=dSdytgVLC8vZ2aF-iTQJ0DnbV2CrzKePLe4X9yWUu40,1285 +cuda/pathfinder/_dynamic_libs/__pycache__/canary_probe_subprocess.cpython-312.pyc,, +cuda/pathfinder/_dynamic_libs/__pycache__/find_nvidia_dynamic_lib.cpython-312.pyc,, +cuda/pathfinder/_dynamic_libs/__pycache__/load_dl_common.cpython-312.pyc,, +cuda/pathfinder/_dynamic_libs/__pycache__/load_dl_linux.cpython-312.pyc,, +cuda/pathfinder/_dynamic_libs/__pycache__/load_dl_windows.cpython-312.pyc,, +cuda/pathfinder/_dynamic_libs/__pycache__/load_nvidia_dynamic_lib.cpython-312.pyc,, +cuda/pathfinder/_dynamic_libs/__pycache__/supported_nvidia_libs.cpython-312.pyc,, +cuda/pathfinder/_dynamic_libs/canary_probe_subprocess.py,sha256=GdzJwYzbtdsOrUB9LsPhUeImapPG0YE4M9o4USBdQxc,994 +cuda/pathfinder/_dynamic_libs/find_nvidia_dynamic_lib.py,sha256=kNVNzDU3mwvXc9O444vdJRbGD09KlS6qIo7dmHSYkhA,10197 +cuda/pathfinder/_dynamic_libs/load_dl_common.py,sha256=vOuh20EKkNKwbRUZFFP621RcjCvgSoKTzaXVmRyW8-k,831 +cuda/pathfinder/_dynamic_libs/load_dl_linux.py,sha256=aSKBgYOm0CCjT-0TixLOgNOtD4kvHHe8gjCBH2aXhjk,8010 +cuda/pathfinder/_dynamic_libs/load_dl_windows.py,sha256=hznwPky7mh5GRYUw12VWcuqrrV-tDLQHCrR6uevbLyQ,6066 +cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py,sha256=Y8ZBeOw3_LfQKcPqTvd3yfE3OO8cIZXPfcg9aFLZwnU,11857 +cuda/pathfinder/_dynamic_libs/supported_nvidia_libs.py,sha256=G_Y4C4OdVu768Y8D018PiqhzC7i4Yhi5if3A_aH7oJ0,14330 +cuda/pathfinder/_headers/__pycache__/find_nvidia_headers.cpython-312.pyc,, +cuda/pathfinder/_headers/__pycache__/supported_nvidia_headers.cpython-312.pyc,, +cuda/pathfinder/_headers/find_nvidia_headers.py,sha256=FUxWFuHe2yZ91bAvt32vJOat0RSyUwNtAstofhCEOrs,7315 +cuda/pathfinder/_headers/supported_nvidia_headers.py,sha256=GIC9Bkh2ufPbOSXxa5vBpsdVnkXCI67F-UCvtdLcYkE,3374 +cuda/pathfinder/_static_libs/__pycache__/find_bitcode_lib.cpython-312.pyc,, +cuda/pathfinder/_static_libs/find_bitcode_lib.py,sha256=W_AyfNCZ9zx8-2W4ZdYjH6aiH6ic2u9ZlDaY-PsvwiU,5460 +cuda/pathfinder/_utils/__pycache__/env_vars.cpython-312.pyc,, +cuda/pathfinder/_utils/__pycache__/find_site_packages_dll.cpython-312.pyc,, +cuda/pathfinder/_utils/__pycache__/find_site_packages_so.cpython-312.pyc,, +cuda/pathfinder/_utils/__pycache__/find_sub_dirs.cpython-312.pyc,, +cuda/pathfinder/_utils/__pycache__/platform_aware.cpython-312.pyc,, +cuda/pathfinder/_utils/__pycache__/spawned_process_runner.cpython-312.pyc,, +cuda/pathfinder/_utils/env_vars.py,sha256=wRpf6XnohhKIpiGz_HbvnGg9C9u0uY5uU3t9OyYCl1A,1756 +cuda/pathfinder/_utils/find_site_packages_dll.py,sha256=HixAy-Ubd5Gu9DZwCLAu2jvlEr1nfolOSZ_7GX3ZgbM,1143 +cuda/pathfinder/_utils/find_site_packages_so.py,sha256=_udOAI8RtTXRJWz-X6TUy-g9W_nW8PeAcG61k1AVSLw,1638 +cuda/pathfinder/_utils/find_sub_dirs.py,sha256=L5SIXikH8U-6BCIvTotR4xwAelWUvJohgZ1RAuxg2eA,2032 +cuda/pathfinder/_utils/platform_aware.py,sha256=oOCqy_VyNdEdISVM2ilde8Kl6rOeElItvsAnw0WD9b8,463 +cuda/pathfinder/_utils/spawned_process_runner.py,sha256=De1idhZq6v3JF2Ar2d5e08XBsoqXyJiuPZlXilYi7Cc,4355 +cuda/pathfinder/_version.py,sha256=cWq_fKDGb85xfXQfl5uEbSwbhongIFne__ApDDkK9B4,712 +cuda_pathfinder-1.4.0.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +cuda_pathfinder-1.4.0.dist-info/METADATA,sha256=XrzGu-7KW6TGCSXmdhPRWx9t_nxpF4xfnS32vcuqG60,1906 +cuda_pathfinder-1.4.0.dist-info/RECORD,, +cuda_pathfinder-1.4.0.dist-info/WHEEL,sha256=YCfwYGOYMi5Jhw2fU4yNgwErybb2IX5PEwBKV4ZbdBo,91 +cuda_pathfinder-1.4.0.dist-info/licenses/LICENSE,sha256=DVQuDIgE45qn836wDaWnYhSdxoLXgpRRKH4RuTjpRZQ,10174 +cuda_pathfinder-1.4.0.dist-info/top_level.txt,sha256=U5vpnvwNpaJF8bl4KnoUluMDTRt0J972FipwXjgNQ3A,5 diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/cuda_pathfinder-1.4.0.dist-info/WHEEL b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/cuda_pathfinder-1.4.0.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..1ef5583317a5e59140e3f1c85c2db91aec0961e8 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/cuda_pathfinder-1.4.0.dist-info/WHEEL @@ -0,0 +1,5 @@ +Wheel-Version: 1.0 +Generator: setuptools (82.0.0) +Root-Is-Purelib: true +Tag: py3-none-any + diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/cuda_pathfinder-1.4.0.dist-info/top_level.txt b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/cuda_pathfinder-1.4.0.dist-info/top_level.txt new file mode 100644 index 0000000000000000000000000000000000000000..089f044f442ab95aabf9de9eb52b63efbac9f7b9 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/cuda_pathfinder-1.4.0.dist-info/top_level.txt @@ -0,0 +1 @@ +cuda diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/fsspec/caching.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/fsspec/caching.py new file mode 100644 index 0000000000000000000000000000000000000000..fb76aa13cb1c6164af5ef938355fddd5b87397e5 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/fsspec/caching.py @@ -0,0 +1,950 @@ +from __future__ import annotations + +import collections +import functools +import logging +import math +import os +import threading +import warnings +from concurrent.futures import Future, ThreadPoolExecutor +from typing import ( + TYPE_CHECKING, + Any, + Callable, + ClassVar, + Generic, + NamedTuple, + OrderedDict, + TypeVar, +) + +if TYPE_CHECKING: + import mmap + + from typing_extensions import ParamSpec + + P = ParamSpec("P") +else: + P = TypeVar("P") + +T = TypeVar("T") + + +logger = logging.getLogger("fsspec") + +Fetcher = Callable[[int, int], bytes] # Maps (start, end) to bytes + + +class BaseCache: + """Pass-though cache: doesn't keep anything, calls every time + + Acts as base class for other cachers + + Parameters + ---------- + blocksize: int + How far to read ahead in numbers of bytes + fetcher: func + Function of the form f(start, end) which gets bytes from remote as + specified + size: int + How big this file is + """ + + name: ClassVar[str] = "none" + + def __init__(self, blocksize: int, fetcher: Fetcher, size: int) -> None: + self.blocksize = blocksize + self.nblocks = 0 + self.fetcher = fetcher + self.size = size + self.hit_count = 0 + self.miss_count = 0 + # the bytes that we actually requested + self.total_requested_bytes = 0 + + def _fetch(self, start: int | None, stop: int | None) -> bytes: + if start is None: + start = 0 + if stop is None: + stop = self.size + if start >= self.size or start >= stop: + return b"" + return self.fetcher(start, stop) + + def _reset_stats(self) -> None: + """Reset hit and miss counts for a more ganular report e.g. by file.""" + self.hit_count = 0 + self.miss_count = 0 + self.total_requested_bytes = 0 + + def _log_stats(self) -> str: + """Return a formatted string of the cache statistics.""" + if self.hit_count == 0 and self.miss_count == 0: + # a cache that does nothing, this is for logs only + return "" + return " , %s: %d hits, %d misses, %d total requested bytes" % ( + self.name, + self.hit_count, + self.miss_count, + self.total_requested_bytes, + ) + + def __repr__(self) -> str: + # TODO: use rich for better formatting + return f""" + <{self.__class__.__name__}: + block size : {self.blocksize} + block count : {self.nblocks} + file size : {self.size} + cache hits : {self.hit_count} + cache misses: {self.miss_count} + total requested bytes: {self.total_requested_bytes}> + """ + + +class MMapCache(BaseCache): + """memory-mapped sparse file cache + + Opens temporary file, which is filled blocks-wise when data is requested. + Ensure there is enough disc space in the temporary location. + + This cache method might only work on posix + """ + + name = "mmap" + + def __init__( + self, + blocksize: int, + fetcher: Fetcher, + size: int, + location: str | None = None, + blocks: set[int] | None = None, + ) -> None: + super().__init__(blocksize, fetcher, size) + self.blocks = set() if blocks is None else blocks + self.location = location + self.cache = self._makefile() + + def _makefile(self) -> mmap.mmap | bytearray: + import mmap + import tempfile + + if self.size == 0: + return bytearray() + + # posix version + if self.location is None or not os.path.exists(self.location): + if self.location is None: + fd = tempfile.TemporaryFile() + self.blocks = set() + else: + fd = open(self.location, "wb+") + fd.seek(self.size - 1) + fd.write(b"1") + fd.flush() + else: + fd = open(self.location, "r+b") + + return mmap.mmap(fd.fileno(), self.size) + + def _fetch(self, start: int | None, end: int | None) -> bytes: + logger.debug(f"MMap cache fetching {start}-{end}") + if start is None: + start = 0 + if end is None: + end = self.size + if start >= self.size or start >= end: + return b"" + start_block = start // self.blocksize + end_block = end // self.blocksize + need = [i for i in range(start_block, end_block + 1) if i not in self.blocks] + hits = [i for i in range(start_block, end_block + 1) if i in self.blocks] + self.miss_count += len(need) + self.hit_count += len(hits) + while need: + # TODO: not a for loop so we can consolidate blocks later to + # make fewer fetch calls; this could be parallel + i = need.pop(0) + + sstart = i * self.blocksize + send = min(sstart + self.blocksize, self.size) + self.total_requested_bytes += send - sstart + logger.debug(f"MMap get block #{i} ({sstart}-{send})") + self.cache[sstart:send] = self.fetcher(sstart, send) + self.blocks.add(i) + + return self.cache[start:end] + + def __getstate__(self) -> dict[str, Any]: + state = self.__dict__.copy() + # Remove the unpicklable entries. + del state["cache"] + return state + + def __setstate__(self, state: dict[str, Any]) -> None: + # Restore instance attributes + self.__dict__.update(state) + self.cache = self._makefile() + + +class ReadAheadCache(BaseCache): + """Cache which reads only when we get beyond a block of data + + This is a much simpler version of BytesCache, and does not attempt to + fill holes in the cache or keep fragments alive. It is best suited to + many small reads in a sequential order (e.g., reading lines from a file). + """ + + name = "readahead" + + def __init__(self, blocksize: int, fetcher: Fetcher, size: int) -> None: + super().__init__(blocksize, fetcher, size) + self.cache = b"" + self.start = 0 + self.end = 0 + + def _fetch(self, start: int | None, end: int | None) -> bytes: + if start is None: + start = 0 + if end is None or end > self.size: + end = self.size + if start >= self.size or start >= end: + return b"" + l = end - start + if start >= self.start and end <= self.end: + # cache hit + self.hit_count += 1 + return self.cache[start - self.start : end - self.start] + elif self.start <= start < self.end: + # partial hit + self.miss_count += 1 + part = self.cache[start - self.start :] + l -= len(part) + start = self.end + else: + # miss + self.miss_count += 1 + part = b"" + end = min(self.size, end + self.blocksize) + self.total_requested_bytes += end - start + self.cache = self.fetcher(start, end) # new block replaces old + self.start = start + self.end = self.start + len(self.cache) + return part + self.cache[:l] + + +class FirstChunkCache(BaseCache): + """Caches the first block of a file only + + This may be useful for file types where the metadata is stored in the header, + but is randomly accessed. + """ + + name = "first" + + def __init__(self, blocksize: int, fetcher: Fetcher, size: int) -> None: + if blocksize > size: + # this will buffer the whole thing + blocksize = size + super().__init__(blocksize, fetcher, size) + self.cache: bytes | None = None + + def _fetch(self, start: int | None, end: int | None) -> bytes: + start = start or 0 + if start > self.size: + logger.debug("FirstChunkCache: requested start > file size") + return b"" + + end = min(end, self.size) + + if start < self.blocksize: + if self.cache is None: + self.miss_count += 1 + if end > self.blocksize: + self.total_requested_bytes += end + data = self.fetcher(0, end) + self.cache = data[: self.blocksize] + return data[start:] + self.cache = self.fetcher(0, self.blocksize) + self.total_requested_bytes += self.blocksize + part = self.cache[start:end] + if end > self.blocksize: + self.total_requested_bytes += end - self.blocksize + part += self.fetcher(self.blocksize, end) + self.hit_count += 1 + return part + else: + self.miss_count += 1 + self.total_requested_bytes += end - start + return self.fetcher(start, end) + + +class BlockCache(BaseCache): + """ + Cache holding memory as a set of blocks. + + Requests are only ever made ``blocksize`` at a time, and are + stored in an LRU cache. The least recently accessed block is + discarded when more than ``maxblocks`` are stored. + + Parameters + ---------- + blocksize : int + The number of bytes to store in each block. + Requests are only ever made for ``blocksize``, so this + should balance the overhead of making a request against + the granularity of the blocks. + fetcher : Callable + size : int + The total size of the file being cached. + maxblocks : int + The maximum number of blocks to cache for. The maximum memory + use for this cache is then ``blocksize * maxblocks``. + """ + + name = "blockcache" + + def __init__( + self, blocksize: int, fetcher: Fetcher, size: int, maxblocks: int = 32 + ) -> None: + super().__init__(blocksize, fetcher, size) + self.nblocks = math.ceil(size / blocksize) + self.maxblocks = maxblocks + self._fetch_block_cached = functools.lru_cache(maxblocks)(self._fetch_block) + + def cache_info(self): + """ + The statistics on the block cache. + + Returns + ------- + NamedTuple + Returned directly from the LRU Cache used internally. + """ + return self._fetch_block_cached.cache_info() + + def __getstate__(self) -> dict[str, Any]: + state = self.__dict__ + del state["_fetch_block_cached"] + return state + + def __setstate__(self, state: dict[str, Any]) -> None: + self.__dict__.update(state) + self._fetch_block_cached = functools.lru_cache(state["maxblocks"])( + self._fetch_block + ) + + def _fetch(self, start: int | None, end: int | None) -> bytes: + if start is None: + start = 0 + if end is None: + end = self.size + if start >= self.size or start >= end: + return b"" + + # byte position -> block numbers + start_block_number = start // self.blocksize + end_block_number = end // self.blocksize + + # these are cached, so safe to do multiple calls for the same start and end. + for block_number in range(start_block_number, end_block_number + 1): + self._fetch_block_cached(block_number) + + return self._read_cache( + start, + end, + start_block_number=start_block_number, + end_block_number=end_block_number, + ) + + def _fetch_block(self, block_number: int) -> bytes: + """ + Fetch the block of data for `block_number`. + """ + if block_number > self.nblocks: + raise ValueError( + f"'block_number={block_number}' is greater than " + f"the number of blocks ({self.nblocks})" + ) + + start = block_number * self.blocksize + end = start + self.blocksize + self.total_requested_bytes += end - start + self.miss_count += 1 + logger.info("BlockCache fetching block %d", block_number) + block_contents = super()._fetch(start, end) + return block_contents + + def _read_cache( + self, start: int, end: int, start_block_number: int, end_block_number: int + ) -> bytes: + """ + Read from our block cache. + + Parameters + ---------- + start, end : int + The start and end byte positions. + start_block_number, end_block_number : int + The start and end block numbers. + """ + start_pos = start % self.blocksize + end_pos = end % self.blocksize + + self.hit_count += 1 + if start_block_number == end_block_number: + block: bytes = self._fetch_block_cached(start_block_number) + return block[start_pos:end_pos] + + else: + # read from the initial + out = [self._fetch_block_cached(start_block_number)[start_pos:]] + + # intermediate blocks + # Note: it'd be nice to combine these into one big request. However + # that doesn't play nicely with our LRU cache. + out.extend( + map( + self._fetch_block_cached, + range(start_block_number + 1, end_block_number), + ) + ) + + # final block + out.append(self._fetch_block_cached(end_block_number)[:end_pos]) + + return b"".join(out) + + +class BytesCache(BaseCache): + """Cache which holds data in a in-memory bytes object + + Implements read-ahead by the block size, for semi-random reads progressing + through the file. + + Parameters + ---------- + trim: bool + As we read more data, whether to discard the start of the buffer when + we are more than a blocksize ahead of it. + """ + + name: ClassVar[str] = "bytes" + + def __init__( + self, blocksize: int, fetcher: Fetcher, size: int, trim: bool = True + ) -> None: + super().__init__(blocksize, fetcher, size) + self.cache = b"" + self.start: int | None = None + self.end: int | None = None + self.trim = trim + + def _fetch(self, start: int | None, end: int | None) -> bytes: + # TODO: only set start/end after fetch, in case it fails? + # is this where retry logic might go? + if start is None: + start = 0 + if end is None: + end = self.size + if start >= self.size or start >= end: + return b"" + if ( + self.start is not None + and start >= self.start + and self.end is not None + and end < self.end + ): + # cache hit: we have all the required data + offset = start - self.start + self.hit_count += 1 + return self.cache[offset : offset + end - start] + + if self.blocksize: + bend = min(self.size, end + self.blocksize) + else: + bend = end + + if bend == start or start > self.size: + return b"" + + if (self.start is None or start < self.start) and ( + self.end is None or end > self.end + ): + # First read, or extending both before and after + self.total_requested_bytes += bend - start + self.miss_count += 1 + self.cache = self.fetcher(start, bend) + self.start = start + else: + assert self.start is not None + assert self.end is not None + self.miss_count += 1 + + if start < self.start: + if self.end is None or self.end - end > self.blocksize: + self.total_requested_bytes += bend - start + self.cache = self.fetcher(start, bend) + self.start = start + else: + self.total_requested_bytes += self.start - start + new = self.fetcher(start, self.start) + self.start = start + self.cache = new + self.cache + elif self.end is not None and bend > self.end: + if self.end > self.size: + pass + elif end - self.end > self.blocksize: + self.total_requested_bytes += bend - start + self.cache = self.fetcher(start, bend) + self.start = start + else: + self.total_requested_bytes += bend - self.end + new = self.fetcher(self.end, bend) + self.cache = self.cache + new + + self.end = self.start + len(self.cache) + offset = start - self.start + out = self.cache[offset : offset + end - start] + if self.trim: + num = (self.end - self.start) // (self.blocksize + 1) + if num > 1: + self.start += self.blocksize * num + self.cache = self.cache[self.blocksize * num :] + return out + + def __len__(self) -> int: + return len(self.cache) + + +class AllBytes(BaseCache): + """Cache entire contents of the file""" + + name: ClassVar[str] = "all" + + def __init__( + self, + blocksize: int | None = None, + fetcher: Fetcher | None = None, + size: int | None = None, + data: bytes | None = None, + ) -> None: + super().__init__(blocksize, fetcher, size) # type: ignore[arg-type] + if data is None: + self.miss_count += 1 + self.total_requested_bytes += self.size + data = self.fetcher(0, self.size) + self.data = data + + def _fetch(self, start: int | None, stop: int | None) -> bytes: + self.hit_count += 1 + return self.data[start:stop] + + +class KnownPartsOfAFile(BaseCache): + """ + Cache holding known file parts. + + Parameters + ---------- + blocksize: int + How far to read ahead in numbers of bytes + fetcher: func + Function of the form f(start, end) which gets bytes from remote as + specified + size: int + How big this file is + data: dict + A dictionary mapping explicit `(start, stop)` file-offset tuples + with known bytes. + strict: bool, default True + Whether to fetch reads that go beyond a known byte-range boundary. + If `False`, any read that ends outside a known part will be zero + padded. Note that zero padding will not be used for reads that + begin outside a known byte-range. + """ + + name: ClassVar[str] = "parts" + + def __init__( + self, + blocksize: int, + fetcher: Fetcher, + size: int, + data: dict[tuple[int, int], bytes] = {}, + strict: bool = True, + **_: Any, + ): + super().__init__(blocksize, fetcher, size) + self.strict = strict + + # simple consolidation of contiguous blocks + if data: + old_offsets = sorted(data.keys()) + offsets = [old_offsets[0]] + blocks = [data.pop(old_offsets[0])] + for start, stop in old_offsets[1:]: + start0, stop0 = offsets[-1] + if start == stop0: + offsets[-1] = (start0, stop) + blocks[-1] += data.pop((start, stop)) + else: + offsets.append((start, stop)) + blocks.append(data.pop((start, stop))) + + self.data = dict(zip(offsets, blocks)) + else: + self.data = data + + def _fetch(self, start: int | None, stop: int | None) -> bytes: + if start is None: + start = 0 + if stop is None: + stop = self.size + + out = b"" + for (loc0, loc1), data in self.data.items(): + # If self.strict=False, use zero-padded data + # for reads beyond the end of a "known" buffer + if loc0 <= start < loc1: + off = start - loc0 + out = data[off : off + stop - start] + if not self.strict or loc0 <= stop <= loc1: + # The request is within a known range, or + # it begins within a known range, and we + # are allowed to pad reads beyond the + # buffer with zero + out += b"\x00" * (stop - start - len(out)) + self.hit_count += 1 + return out + else: + # The request ends outside a known range, + # and we are being "strict" about reads + # beyond the buffer + start = loc1 + break + + # We only get here if there is a request outside the + # known parts of the file. In an ideal world, this + # should never happen + if self.fetcher is None: + # We cannot fetch the data, so raise an error + raise ValueError(f"Read is outside the known file parts: {(start, stop)}. ") + # We can fetch the data, but should warn the user + # that this may be slow + warnings.warn( + f"Read is outside the known file parts: {(start, stop)}. " + f"IO/caching performance may be poor!" + ) + logger.debug(f"KnownPartsOfAFile cache fetching {start}-{stop}") + self.total_requested_bytes += stop - start + self.miss_count += 1 + return out + super()._fetch(start, stop) + + +class UpdatableLRU(Generic[P, T]): + """ + Custom implementation of LRU cache that allows updating keys + + Used by BackgroudBlockCache + """ + + class CacheInfo(NamedTuple): + hits: int + misses: int + maxsize: int + currsize: int + + def __init__(self, func: Callable[P, T], max_size: int = 128) -> None: + self._cache: OrderedDict[Any, T] = collections.OrderedDict() + self._func = func + self._max_size = max_size + self._hits = 0 + self._misses = 0 + self._lock = threading.Lock() + + def __call__(self, *args: P.args, **kwargs: P.kwargs) -> T: + if kwargs: + raise TypeError(f"Got unexpected keyword argument {kwargs.keys()}") + with self._lock: + if args in self._cache: + self._cache.move_to_end(args) + self._hits += 1 + return self._cache[args] + + result = self._func(*args, **kwargs) + + with self._lock: + self._cache[args] = result + self._misses += 1 + if len(self._cache) > self._max_size: + self._cache.popitem(last=False) + + return result + + def is_key_cached(self, *args: Any) -> bool: + with self._lock: + return args in self._cache + + def add_key(self, result: T, *args: Any) -> None: + with self._lock: + self._cache[args] = result + if len(self._cache) > self._max_size: + self._cache.popitem(last=False) + + def cache_info(self) -> UpdatableLRU.CacheInfo: + with self._lock: + return self.CacheInfo( + maxsize=self._max_size, + currsize=len(self._cache), + hits=self._hits, + misses=self._misses, + ) + + +class BackgroundBlockCache(BaseCache): + """ + Cache holding memory as a set of blocks with pre-loading of + the next block in the background. + + Requests are only ever made ``blocksize`` at a time, and are + stored in an LRU cache. The least recently accessed block is + discarded when more than ``maxblocks`` are stored. If the + next block is not in cache, it is loaded in a separate thread + in non-blocking way. + + Parameters + ---------- + blocksize : int + The number of bytes to store in each block. + Requests are only ever made for ``blocksize``, so this + should balance the overhead of making a request against + the granularity of the blocks. + fetcher : Callable + size : int + The total size of the file being cached. + maxblocks : int + The maximum number of blocks to cache for. The maximum memory + use for this cache is then ``blocksize * maxblocks``. + """ + + name: ClassVar[str] = "background" + + def __init__( + self, blocksize: int, fetcher: Fetcher, size: int, maxblocks: int = 32 + ) -> None: + super().__init__(blocksize, fetcher, size) + self.nblocks = math.ceil(size / blocksize) + self.maxblocks = maxblocks + self._fetch_block_cached = UpdatableLRU(self._fetch_block, maxblocks) + + self._thread_executor = ThreadPoolExecutor(max_workers=1) + self._fetch_future_block_number: int | None = None + self._fetch_future: Future[bytes] | None = None + self._fetch_future_lock = threading.Lock() + + def cache_info(self) -> UpdatableLRU.CacheInfo: + """ + The statistics on the block cache. + + Returns + ------- + NamedTuple + Returned directly from the LRU Cache used internally. + """ + return self._fetch_block_cached.cache_info() + + def __getstate__(self) -> dict[str, Any]: + state = self.__dict__ + del state["_fetch_block_cached"] + del state["_thread_executor"] + del state["_fetch_future_block_number"] + del state["_fetch_future"] + del state["_fetch_future_lock"] + return state + + def __setstate__(self, state) -> None: + self.__dict__.update(state) + self._fetch_block_cached = UpdatableLRU(self._fetch_block, state["maxblocks"]) + self._thread_executor = ThreadPoolExecutor(max_workers=1) + self._fetch_future_block_number = None + self._fetch_future = None + self._fetch_future_lock = threading.Lock() + + def _fetch(self, start: int | None, end: int | None) -> bytes: + if start is None: + start = 0 + if end is None: + end = self.size + if start >= self.size or start >= end: + return b"" + + # byte position -> block numbers + start_block_number = start // self.blocksize + end_block_number = end // self.blocksize + + fetch_future_block_number = None + fetch_future = None + with self._fetch_future_lock: + # Background thread is running. Check we we can or must join it. + if self._fetch_future is not None: + assert self._fetch_future_block_number is not None + if self._fetch_future.done(): + logger.info("BlockCache joined background fetch without waiting.") + self._fetch_block_cached.add_key( + self._fetch_future.result(), self._fetch_future_block_number + ) + # Cleanup the fetch variables. Done with fetching the block. + self._fetch_future_block_number = None + self._fetch_future = None + else: + # Must join if we need the block for the current fetch + must_join = bool( + start_block_number + <= self._fetch_future_block_number + <= end_block_number + ) + if must_join: + # Copy to the local variables to release lock + # before waiting for result + fetch_future_block_number = self._fetch_future_block_number + fetch_future = self._fetch_future + + # Cleanup the fetch variables. Have a local copy. + self._fetch_future_block_number = None + self._fetch_future = None + + # Need to wait for the future for the current read + if fetch_future is not None: + logger.info("BlockCache waiting for background fetch.") + # Wait until result and put it in cache + self._fetch_block_cached.add_key( + fetch_future.result(), fetch_future_block_number + ) + + # these are cached, so safe to do multiple calls for the same start and end. + for block_number in range(start_block_number, end_block_number + 1): + self._fetch_block_cached(block_number) + + # fetch next block in the background if nothing is running in the background, + # the block is within file and it is not already cached + end_block_plus_1 = end_block_number + 1 + with self._fetch_future_lock: + if ( + self._fetch_future is None + and end_block_plus_1 <= self.nblocks + and not self._fetch_block_cached.is_key_cached(end_block_plus_1) + ): + self._fetch_future_block_number = end_block_plus_1 + self._fetch_future = self._thread_executor.submit( + self._fetch_block, end_block_plus_1, "async" + ) + + return self._read_cache( + start, + end, + start_block_number=start_block_number, + end_block_number=end_block_number, + ) + + def _fetch_block(self, block_number: int, log_info: str = "sync") -> bytes: + """ + Fetch the block of data for `block_number`. + """ + if block_number > self.nblocks: + raise ValueError( + f"'block_number={block_number}' is greater than " + f"the number of blocks ({self.nblocks})" + ) + + start = block_number * self.blocksize + end = start + self.blocksize + logger.info("BlockCache fetching block (%s) %d", log_info, block_number) + self.total_requested_bytes += end - start + self.miss_count += 1 + block_contents = super()._fetch(start, end) + return block_contents + + def _read_cache( + self, start: int, end: int, start_block_number: int, end_block_number: int + ) -> bytes: + """ + Read from our block cache. + + Parameters + ---------- + start, end : int + The start and end byte positions. + start_block_number, end_block_number : int + The start and end block numbers. + """ + start_pos = start % self.blocksize + end_pos = end % self.blocksize + + # kind of pointless to count this as a hit, but it is + self.hit_count += 1 + + if start_block_number == end_block_number: + block = self._fetch_block_cached(start_block_number) + return block[start_pos:end_pos] + + else: + # read from the initial + out = [self._fetch_block_cached(start_block_number)[start_pos:]] + + # intermediate blocks + # Note: it'd be nice to combine these into one big request. However + # that doesn't play nicely with our LRU cache. + out.extend( + map( + self._fetch_block_cached, + range(start_block_number + 1, end_block_number), + ) + ) + + # final block + out.append(self._fetch_block_cached(end_block_number)[:end_pos]) + + return b"".join(out) + + +caches: dict[str | None, type[BaseCache]] = { + # one custom case + None: BaseCache, +} + + +def register_cache(cls: type[BaseCache], clobber: bool = False) -> None: + """'Register' cache implementation. + + Parameters + ---------- + clobber: bool, optional + If set to True (default is False) - allow to overwrite existing + entry. + + Raises + ------ + ValueError + """ + name = cls.name + if not clobber and name in caches: + raise ValueError(f"Cache with name {name!r} is already known: {caches[name]}") + caches[name] = cls + + +for c in ( + BaseCache, + MMapCache, + BytesCache, + ReadAheadCache, + BlockCache, + FirstChunkCache, + AllBytes, + KnownPartsOfAFile, + BackgroundBlockCache, +): + register_cache(c) diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/fsspec/dircache.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/fsspec/dircache.py new file mode 100644 index 0000000000000000000000000000000000000000..eca19566b135e5a7a4f6e7407d56411ec58bfe44 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/fsspec/dircache.py @@ -0,0 +1,98 @@ +import time +from collections.abc import MutableMapping +from functools import lru_cache + + +class DirCache(MutableMapping): + """ + Caching of directory listings, in a structure like:: + + {"path0": [ + {"name": "path0/file0", + "size": 123, + "type": "file", + ... + }, + {"name": "path0/file1", + }, + ... + ], + "path1": [...] + } + + Parameters to this class control listing expiry or indeed turn + caching off + """ + + def __init__( + self, + use_listings_cache=True, + listings_expiry_time=None, + max_paths=None, + **kwargs, + ): + """ + + Parameters + ---------- + use_listings_cache: bool + If False, this cache never returns items, but always reports KeyError, + and setting items has no effect + listings_expiry_time: int or float (optional) + Time in seconds that a listing is considered valid. If None, + listings do not expire. + max_paths: int (optional) + The number of most recent listings that are considered valid; 'recent' + refers to when the entry was set. + """ + self._cache = {} + self._times = {} + if max_paths: + self._q = lru_cache(max_paths + 1)(lambda key: self._cache.pop(key, None)) + self.use_listings_cache = use_listings_cache + self.listings_expiry_time = listings_expiry_time + self.max_paths = max_paths + + def __getitem__(self, item): + if self.listings_expiry_time is not None: + if self._times.get(item, 0) - time.time() < -self.listings_expiry_time: + del self._cache[item] + if self.max_paths: + self._q(item) + return self._cache[item] # maybe raises KeyError + + def clear(self): + self._cache.clear() + + def __len__(self): + return len(self._cache) + + def __contains__(self, item): + try: + self[item] + return True + except KeyError: + return False + + def __setitem__(self, key, value): + if not self.use_listings_cache: + return + if self.max_paths: + self._q(key) + self._cache[key] = value + if self.listings_expiry_time is not None: + self._times[key] = time.time() + + def __delitem__(self, key): + del self._cache[key] + + def __iter__(self): + entries = list(self._cache) + + return (k for k in entries if k in self) + + def __reduce__(self): + return ( + DirCache, + (self.use_listings_cache, self.listings_expiry_time, self.max_paths), + ) diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/fsspec/gui.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/fsspec/gui.py new file mode 100644 index 0000000000000000000000000000000000000000..113317e5a48c4d5fac062bc84019e97a05e16a0f --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/fsspec/gui.py @@ -0,0 +1,414 @@ +import ast +import contextlib +import logging +import os +import re +from typing import ClassVar, Sequence + +import panel as pn + +from .core import OpenFile, get_filesystem_class, split_protocol +from .registry import known_implementations + +pn.extension() +logger = logging.getLogger("fsspec.gui") + + +class SigSlot: + """Signal-slot mixin, for Panel event passing + + Include this class in a widget manager's superclasses to be able to + register events and callbacks on Panel widgets managed by that class. + + The method ``_register`` should be called as widgets are added, and external + code should call ``connect`` to associate callbacks. + + By default, all signals emit a DEBUG logging statement. + """ + + # names of signals that this class may emit each of which must be + # set by _register for any new instance + signals: ClassVar[Sequence[str]] = [] + # names of actions that this class may respond to + slots: ClassVar[Sequence[str]] = [] + + # each of which must be a method name + + def __init__(self): + self._ignoring_events = False + self._sigs = {} + self._map = {} + self._setup() + + def _setup(self): + """Create GUI elements and register signals""" + self.panel = pn.pane.PaneBase() + # no signals to set up in the base class + + def _register( + self, widget, name, thing="value", log_level=logging.DEBUG, auto=False + ): + """Watch the given attribute of a widget and assign it a named event + + This is normally called at the time a widget is instantiated, in the + class which owns it. + + Parameters + ---------- + widget : pn.layout.Panel or None + Widget to watch. If None, an anonymous signal not associated with + any widget. + name : str + Name of this event + thing : str + Attribute of the given widget to watch + log_level : int + When the signal is triggered, a logging event of the given level + will be fired in the dfviz logger. + auto : bool + If True, automatically connects with a method in this class of the + same name. + """ + if name not in self.signals: + raise ValueError(f"Attempt to assign an undeclared signal: {name}") + self._sigs[name] = { + "widget": widget, + "callbacks": [], + "thing": thing, + "log": log_level, + } + wn = "-".join( + [ + getattr(widget, "name", str(widget)) if widget is not None else "none", + thing, + ] + ) + self._map[wn] = name + if widget is not None: + widget.param.watch(self._signal, thing, onlychanged=True) + if auto and hasattr(self, name): + self.connect(name, getattr(self, name)) + + def _repr_mimebundle_(self, *args, **kwargs): + """Display in a notebook or a server""" + try: + return self.panel._repr_mimebundle_(*args, **kwargs) + except (ValueError, AttributeError): + raise NotImplementedError("Panel does not seem to be set up properly") + + def connect(self, signal, slot): + """Associate call back with given event + + The callback must be a function which takes the "new" value of the + watched attribute as the only parameter. If the callback return False, + this cancels any further processing of the given event. + + Alternatively, the callback can be a string, in which case it means + emitting the correspondingly-named event (i.e., connect to self) + """ + self._sigs[signal]["callbacks"].append(slot) + + def _signal(self, event): + """This is called by a an action on a widget + + Within an self.ignore_events context, nothing happens. + + Tests can execute this method by directly changing the values of + widget components. + """ + if not self._ignoring_events: + wn = "-".join([event.obj.name, event.name]) + if wn in self._map and self._map[wn] in self._sigs: + self._emit(self._map[wn], event.new) + + @contextlib.contextmanager + def ignore_events(self): + """Temporarily turn off events processing in this instance + + (does not propagate to children) + """ + self._ignoring_events = True + try: + yield + finally: + self._ignoring_events = False + + def _emit(self, sig, value=None): + """An event happened, call its callbacks + + This method can be used in tests to simulate message passing without + directly changing visual elements. + + Calling of callbacks will halt whenever one returns False. + """ + logger.log(self._sigs[sig]["log"], f"{sig}: {value}") + for callback in self._sigs[sig]["callbacks"]: + if isinstance(callback, str): + self._emit(callback) + else: + try: + # running callbacks should not break the interface + ret = callback(value) + if ret is False: + break + except Exception as e: + logger.exception( + "Exception (%s) while executing callback for signal: %s", + e, + sig, + ) + + def show(self, threads=False): + """Open a new browser tab and display this instance's interface""" + self.panel.show(threads=threads, verbose=False) + return self + + +class SingleSelect(SigSlot): + """A multiselect which only allows you to select one item for an event""" + + signals = ["_selected", "selected"] # the first is internal + slots = ["set_options", "set_selection", "add", "clear", "select"] + + def __init__(self, **kwargs): + self.kwargs = kwargs + super().__init__() + + def _setup(self): + self.panel = pn.widgets.MultiSelect(**self.kwargs) + self._register(self.panel, "_selected", "value") + self._register(None, "selected") + self.connect("_selected", self.select_one) + + def _signal(self, *args, **kwargs): + super()._signal(*args, **kwargs) + + def select_one(self, *_): + with self.ignore_events(): + val = [self.panel.value[-1]] if self.panel.value else [] + self.panel.value = val + self._emit("selected", self.panel.value) + + def set_options(self, options): + self.panel.options = options + + def clear(self): + self.panel.options = [] + + @property + def value(self): + return self.panel.value + + def set_selection(self, selection): + self.panel.value = [selection] + + +class FileSelector(SigSlot): + """Panel-based graphical file selector widget + + Instances of this widget are interactive and can be displayed in jupyter by having + them as the output of a cell, or in a separate browser tab using ``.show()``. + """ + + signals = [ + "protocol_changed", + "selection_changed", + "directory_entered", + "home_clicked", + "up_clicked", + "go_clicked", + "filters_changed", + ] + slots = ["set_filters", "go_home"] + + def __init__(self, url=None, filters=None, ignore=None, kwargs=None): + """ + + Parameters + ---------- + url : str (optional) + Initial value of the URL to populate the dialog; should include protocol + filters : list(str) (optional) + File endings to include in the listings. If not included, all files are + allowed. Does not affect directories. + If given, the endings will appear as checkboxes in the interface + ignore : list(str) (optional) + Regex(s) of file basename patterns to ignore, e.g., "\\." for typical + hidden files on posix + kwargs : dict (optional) + To pass to file system instance + """ + if url: + self.init_protocol, url = split_protocol(url) + else: + self.init_protocol, url = "file", os.getcwd() + self.init_url = url + self.init_kwargs = (kwargs if isinstance(kwargs, str) else str(kwargs)) or "{}" + self.filters = filters + self.ignore = [re.compile(i) for i in ignore or []] + self._fs = None + super().__init__() + + def _setup(self): + self.url = pn.widgets.TextInput( + name="url", + value=self.init_url, + align="end", + sizing_mode="stretch_width", + width_policy="max", + ) + self.protocol = pn.widgets.Select( + options=sorted(known_implementations), + value=self.init_protocol, + name="protocol", + align="center", + ) + self.kwargs = pn.widgets.TextInput( + name="kwargs", value=self.init_kwargs, align="center" + ) + self.go = pn.widgets.Button(name="⇨", align="end", width=45) + self.main = SingleSelect(size=10) + self.home = pn.widgets.Button(name="🏠", width=40, height=30, align="end") + self.up = pn.widgets.Button(name="‹", width=30, height=30, align="end") + + self._register(self.protocol, "protocol_changed", auto=True) + self._register(self.go, "go_clicked", "clicks", auto=True) + self._register(self.up, "up_clicked", "clicks", auto=True) + self._register(self.home, "home_clicked", "clicks", auto=True) + self._register(None, "selection_changed") + self.main.connect("selected", self.selection_changed) + self._register(None, "directory_entered") + self.prev_protocol = self.protocol.value + self.prev_kwargs = self.storage_options + + self.filter_sel = pn.widgets.CheckBoxGroup( + value=[], options=[], inline=False, align="end", width_policy="min" + ) + self._register(self.filter_sel, "filters_changed", auto=True) + + self.panel = pn.Column( + pn.Row(self.protocol, self.kwargs), + pn.Row(self.home, self.up, self.url, self.go, self.filter_sel), + self.main.panel, + ) + self.set_filters(self.filters) + self.go_clicked() + + def set_filters(self, filters=None): + self.filters = filters + if filters: + self.filter_sel.options = filters + self.filter_sel.value = filters + else: + self.filter_sel.options = [] + self.filter_sel.value = [] + + @property + def storage_options(self): + """Value of the kwargs box as a dictionary""" + return ast.literal_eval(self.kwargs.value) or {} + + @property + def fs(self): + """Current filesystem instance""" + if self._fs is None: + cls = get_filesystem_class(self.protocol.value) + self._fs = cls(**self.storage_options) + return self._fs + + @property + def urlpath(self): + """URL of currently selected item""" + return ( + (f"{self.protocol.value}://{self.main.value[0]}") + if self.main.value + else None + ) + + def open_file(self, mode="rb", compression=None, encoding=None): + """Create OpenFile instance for the currently selected item + + For example, in a notebook you might do something like + + .. code-block:: + + [ ]: sel = FileSelector(); sel + + # user selects their file + + [ ]: with sel.open_file('rb') as f: + ... out = f.read() + + Parameters + ---------- + mode: str (optional) + Open mode for the file. + compression: str (optional) + The interact with the file as compressed. Set to 'infer' to guess + compression from the file ending + encoding: str (optional) + If using text mode, use this encoding; defaults to UTF8. + """ + if self.urlpath is None: + raise ValueError("No file selected") + return OpenFile(self.fs, self.urlpath, mode, compression, encoding) + + def filters_changed(self, values): + self.filters = values + self.go_clicked() + + def selection_changed(self, *_): + if self.urlpath is None: + return + if self.fs.isdir(self.urlpath): + self.url.value = self.fs._strip_protocol(self.urlpath) + self.go_clicked() + + def go_clicked(self, *_): + if ( + self.prev_protocol != self.protocol.value + or self.prev_kwargs != self.storage_options + ): + self._fs = None # causes fs to be recreated + self.prev_protocol = self.protocol.value + self.prev_kwargs = self.storage_options + listing = sorted( + self.fs.ls(self.url.value, detail=True), key=lambda x: x["name"] + ) + listing = [ + l + for l in listing + if not any(i.match(l["name"].rsplit("/", 1)[-1]) for i in self.ignore) + ] + folders = { + "📁 " + o["name"].rsplit("/", 1)[-1]: o["name"] + for o in listing + if o["type"] == "directory" + } + files = { + "📄 " + o["name"].rsplit("/", 1)[-1]: o["name"] + for o in listing + if o["type"] == "file" + } + if self.filters: + files = { + k: v + for k, v in files.items() + if any(v.endswith(ext) for ext in self.filters) + } + self.main.set_options(dict(**folders, **files)) + + def protocol_changed(self, *_): + self._fs = None + self.main.options = [] + self.url.value = "" + + def home_clicked(self, *_): + self.protocol.value = self.init_protocol + self.kwargs.value = self.init_kwargs + self.url.value = self.init_url + self.go_clicked() + + def up_clicked(self, *_): + self.url.value = self.fs._parent(self.url.value) + self.go_clicked() diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/fsspec/spec.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/fsspec/spec.py new file mode 100644 index 0000000000000000000000000000000000000000..bcc01f514e479212b0521ded79949bbbb9541cc3 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/fsspec/spec.py @@ -0,0 +1,1982 @@ +from __future__ import annotations + +import io +import logging +import os +import threading +import warnings +import weakref +from errno import ESPIPE +from glob import has_magic +from hashlib import sha256 +from typing import ClassVar + +from .callbacks import DEFAULT_CALLBACK +from .config import apply_config, conf +from .dircache import DirCache +from .transaction import Transaction +from .utils import ( + _unstrip_protocol, + glob_translate, + isfilelike, + other_paths, + read_block, + stringify_path, + tokenize, +) + +logger = logging.getLogger("fsspec") + + +def make_instance(cls, args, kwargs): + return cls(*args, **kwargs) + + +class _Cached(type): + """ + Metaclass for caching file system instances. + + Notes + ----- + Instances are cached according to + + * The values of the class attributes listed in `_extra_tokenize_attributes` + * The arguments passed to ``__init__``. + + This creates an additional reference to the filesystem, which prevents the + filesystem from being garbage collected when all *user* references go away. + A call to the :meth:`AbstractFileSystem.clear_instance_cache` must *also* + be made for a filesystem instance to be garbage collected. + """ + + def __init__(cls, *args, **kwargs): + super().__init__(*args, **kwargs) + # Note: we intentionally create a reference here, to avoid garbage + # collecting instances when all other references are gone. To really + # delete a FileSystem, the cache must be cleared. + if conf.get("weakref_instance_cache"): # pragma: no cover + # debug option for analysing fork/spawn conditions + cls._cache = weakref.WeakValueDictionary() + else: + cls._cache = {} + cls._pid = os.getpid() + + def __call__(cls, *args, **kwargs): + kwargs = apply_config(cls, kwargs) + extra_tokens = tuple( + getattr(cls, attr, None) for attr in cls._extra_tokenize_attributes + ) + token = tokenize( + cls, cls._pid, threading.get_ident(), *args, *extra_tokens, **kwargs + ) + skip = kwargs.pop("skip_instance_cache", False) + if os.getpid() != cls._pid: + cls._cache.clear() + cls._pid = os.getpid() + if not skip and cls.cachable and token in cls._cache: + cls._latest = token + return cls._cache[token] + else: + obj = super().__call__(*args, **kwargs) + # Setting _fs_token here causes some static linters to complain. + obj._fs_token_ = token + obj.storage_args = args + obj.storage_options = kwargs + if obj.async_impl and obj.mirror_sync_methods: + from .asyn import mirror_sync_methods + + mirror_sync_methods(obj) + + if cls.cachable and not skip: + cls._latest = token + cls._cache[token] = obj + return obj + + +class AbstractFileSystem(metaclass=_Cached): + """ + An abstract super-class for pythonic file-systems + + Implementations are expected to be compatible with or, better, subclass + from here. + """ + + cachable = True # this class can be cached, instances reused + _cached = False + blocksize = 2**22 + sep = "/" + protocol: ClassVar[str | tuple[str, ...]] = "abstract" + _latest = None + async_impl = False + mirror_sync_methods = False + root_marker = "" # For some FSs, may require leading '/' or other character + transaction_type = Transaction + + #: Extra *class attributes* that should be considered when hashing. + _extra_tokenize_attributes = () + + def __init__(self, *args, **storage_options): + """Create and configure file-system instance + + Instances may be cachable, so if similar enough arguments are seen + a new instance is not required. The token attribute exists to allow + implementations to cache instances if they wish. + + A reasonable default should be provided if there are no arguments. + + Subclasses should call this method. + + Parameters + ---------- + use_listings_cache, listings_expiry_time, max_paths: + passed to ``DirCache``, if the implementation supports + directory listing caching. Pass use_listings_cache=False + to disable such caching. + skip_instance_cache: bool + If this is a cachable implementation, pass True here to force + creating a new instance even if a matching instance exists, and prevent + storing this instance. + asynchronous: bool + loop: asyncio-compatible IOLoop or None + """ + if self._cached: + # reusing instance, don't change + return + self._cached = True + self._intrans = False + self._transaction = None + self._invalidated_caches_in_transaction = [] + self.dircache = DirCache(**storage_options) + + if storage_options.pop("add_docs", None): + warnings.warn("add_docs is no longer supported.", FutureWarning) + + if storage_options.pop("add_aliases", None): + warnings.warn("add_aliases has been removed.", FutureWarning) + # This is set in _Cached + self._fs_token_ = None + + @property + def fsid(self): + """Persistent filesystem id that can be used to compare filesystems + across sessions. + """ + raise NotImplementedError + + @property + def _fs_token(self): + return self._fs_token_ + + def __dask_tokenize__(self): + return self._fs_token + + def __hash__(self): + return int(self._fs_token, 16) + + def __eq__(self, other): + return isinstance(other, type(self)) and self._fs_token == other._fs_token + + def __reduce__(self): + return make_instance, (type(self), self.storage_args, self.storage_options) + + @classmethod + def _strip_protocol(cls, path): + """Turn path from fully-qualified to file-system-specific + + May require FS-specific handling, e.g., for relative paths or links. + """ + if isinstance(path, list): + return [cls._strip_protocol(p) for p in path] + path = stringify_path(path) + protos = (cls.protocol,) if isinstance(cls.protocol, str) else cls.protocol + for protocol in protos: + if path.startswith(protocol + "://"): + path = path[len(protocol) + 3 :] + elif path.startswith(protocol + "::"): + path = path[len(protocol) + 2 :] + path = path.rstrip("/") + # use of root_marker to make minimum required path, e.g., "/" + return path or cls.root_marker + + def unstrip_protocol(self, name: str) -> str: + """Format FS-specific path to generic, including protocol""" + protos = (self.protocol,) if isinstance(self.protocol, str) else self.protocol + for protocol in protos: + if name.startswith(f"{protocol}://"): + return name + return f"{protos[0]}://{name}" + + @staticmethod + def _get_kwargs_from_urls(path): + """If kwargs can be encoded in the paths, extract them here + + This should happen before instantiation of the class; incoming paths + then should be amended to strip the options in methods. + + Examples may look like an sftp path "sftp://user@host:/my/path", where + the user and host should become kwargs and later get stripped. + """ + # by default, nothing happens + return {} + + @classmethod + def current(cls): + """Return the most recently instantiated FileSystem + + If no instance has been created, then create one with defaults + """ + if cls._latest in cls._cache: + return cls._cache[cls._latest] + return cls() + + @property + def transaction(self): + """A context within which files are committed together upon exit + + Requires the file class to implement `.commit()` and `.discard()` + for the normal and exception cases. + """ + if self._transaction is None: + self._transaction = self.transaction_type(self) + return self._transaction + + def start_transaction(self): + """Begin write transaction for deferring files, non-context version""" + self._intrans = True + self._transaction = self.transaction_type(self) + return self.transaction + + def end_transaction(self): + """Finish write transaction, non-context version""" + self.transaction.complete() + self._transaction = None + # The invalid cache must be cleared after the transaction is completed. + for path in self._invalidated_caches_in_transaction: + self.invalidate_cache(path) + self._invalidated_caches_in_transaction.clear() + + def invalidate_cache(self, path=None): + """ + Discard any cached directory information + + Parameters + ---------- + path: string or None + If None, clear all listings cached else listings at or under given + path. + """ + # Not necessary to implement invalidation mechanism, may have no cache. + # But if have, you should call this method of parent class from your + # subclass to ensure expiring caches after transacations correctly. + # See the implementation of FTPFileSystem in ftp.py + if self._intrans: + self._invalidated_caches_in_transaction.append(path) + + def mkdir(self, path, create_parents=True, **kwargs): + """ + Create directory entry at path + + For systems that don't have true directories, may create an for + this instance only and not touch the real filesystem + + Parameters + ---------- + path: str + location + create_parents: bool + if True, this is equivalent to ``makedirs`` + kwargs: + may be permissions, etc. + """ + pass # not necessary to implement, may not have directories + + def makedirs(self, path, exist_ok=False): + """Recursively make directories + + Creates directory at path and any intervening required directories. + Raises exception if, for instance, the path already exists but is a + file. + + Parameters + ---------- + path: str + leaf directory name + exist_ok: bool (False) + If False, will error if the target already exists + """ + pass # not necessary to implement, may not have directories + + def rmdir(self, path): + """Remove a directory, if empty""" + pass # not necessary to implement, may not have directories + + def ls(self, path, detail=True, **kwargs): + """List objects at path. + + This should include subdirectories and files at that location. The + difference between a file and a directory must be clear when details + are requested. + + The specific keys, or perhaps a FileInfo class, or similar, is TBD, + but must be consistent across implementations. + Must include: + + - full path to the entry (without protocol) + - size of the entry, in bytes. If the value cannot be determined, will + be ``None``. + - type of entry, "file", "directory" or other + + Additional information + may be present, appropriate to the file-system, e.g., generation, + checksum, etc. + + May use refresh=True|False to allow use of self._ls_from_cache to + check for a saved listing and avoid calling the backend. This would be + common where listing may be expensive. + + Parameters + ---------- + path: str + detail: bool + if True, gives a list of dictionaries, where each is the same as + the result of ``info(path)``. If False, gives a list of paths + (str). + kwargs: may have additional backend-specific options, such as version + information + + Returns + ------- + List of strings if detail is False, or list of directory information + dicts if detail is True. + """ + raise NotImplementedError + + def _ls_from_cache(self, path): + """Check cache for listing + + Returns listing, if found (may be empty list for a directly that exists + but contains nothing), None if not in cache. + """ + parent = self._parent(path) + try: + return self.dircache[path.rstrip("/")] + except KeyError: + pass + try: + files = [ + f + for f in self.dircache[parent] + if f["name"] == path + or (f["name"] == path.rstrip("/") and f["type"] == "directory") + ] + if len(files) == 0: + # parent dir was listed but did not contain this file + raise FileNotFoundError(path) + return files + except KeyError: + pass + + def walk(self, path, maxdepth=None, topdown=True, on_error="omit", **kwargs): + """Return all files belows path + + List all files, recursing into subdirectories; output is iterator-style, + like ``os.walk()``. For a simple list of files, ``find()`` is available. + + When topdown is True, the caller can modify the dirnames list in-place (perhaps + using del or slice assignment), and walk() will + only recurse into the subdirectories whose names remain in dirnames; + this can be used to prune the search, impose a specific order of visiting, + or even to inform walk() about directories the caller creates or renames before + it resumes walk() again. + Modifying dirnames when topdown is False has no effect. (see os.walk) + + Note that the "files" outputted will include anything that is not + a directory, such as links. + + Parameters + ---------- + path: str + Root to recurse into + maxdepth: int + Maximum recursion depth. None means limitless, but not recommended + on link-based file-systems. + topdown: bool (True) + Whether to walk the directory tree from the top downwards or from + the bottom upwards. + on_error: "omit", "raise", a collable + if omit (default), path with exception will simply be empty; + If raise, an underlying exception will be raised; + if callable, it will be called with a single OSError instance as argument + kwargs: passed to ``ls`` + """ + if maxdepth is not None and maxdepth < 1: + raise ValueError("maxdepth must be at least 1") + + path = self._strip_protocol(path) + full_dirs = {} + dirs = {} + files = {} + + detail = kwargs.pop("detail", False) + try: + listing = self.ls(path, detail=True, **kwargs) + except (FileNotFoundError, OSError) as e: + if on_error == "raise": + raise + elif callable(on_error): + on_error(e) + if detail: + return path, {}, {} + return path, [], [] + + for info in listing: + # each info name must be at least [path]/part , but here + # we check also for names like [path]/part/ + pathname = info["name"].rstrip("/") + name = pathname.rsplit("/", 1)[-1] + if info["type"] == "directory" and pathname != path: + # do not include "self" path + full_dirs[name] = pathname + dirs[name] = info + elif pathname == path: + # file-like with same name as give path + files[""] = info + else: + files[name] = info + + if not detail: + dirs = list(dirs) + files = list(files) + + if topdown: + # Yield before recursion if walking top down + yield path, dirs, files + + if maxdepth is not None: + maxdepth -= 1 + if maxdepth < 1: + if not topdown: + yield path, dirs, files + return + + for d in dirs: + yield from self.walk( + full_dirs[d], + maxdepth=maxdepth, + detail=detail, + topdown=topdown, + **kwargs, + ) + + if not topdown: + # Yield after recursion if walking bottom up + yield path, dirs, files + + def find(self, path, maxdepth=None, withdirs=False, detail=False, **kwargs): + """List all files below path. + + Like posix ``find`` command without conditions + + Parameters + ---------- + path : str + maxdepth: int or None + If not None, the maximum number of levels to descend + withdirs: bool + Whether to include directory paths in the output. This is True + when used by glob, but users usually only want files. + kwargs are passed to ``ls``. + """ + # TODO: allow equivalent of -name parameter + path = self._strip_protocol(path) + out = {} + + # Add the root directory if withdirs is requested + # This is needed for posix glob compliance + if withdirs and path != "" and self.isdir(path): + out[path] = self.info(path) + + for _, dirs, files in self.walk(path, maxdepth, detail=True, **kwargs): + if withdirs: + files.update(dirs) + out.update({info["name"]: info for name, info in files.items()}) + if not out and self.isfile(path): + # walk works on directories, but find should also return [path] + # when path happens to be a file + out[path] = {} + names = sorted(out) + if not detail: + return names + else: + return {name: out[name] for name in names} + + def du(self, path, total=True, maxdepth=None, withdirs=False, **kwargs): + """Space used by files and optionally directories within a path + + Directory size does not include the size of its contents. + + Parameters + ---------- + path: str + total: bool + Whether to sum all the file sizes + maxdepth: int or None + Maximum number of directory levels to descend, None for unlimited. + withdirs: bool + Whether to include directory paths in the output. + kwargs: passed to ``find`` + + Returns + ------- + Dict of {path: size} if total=False, or int otherwise, where numbers + refer to bytes used. + """ + sizes = {} + if withdirs and self.isdir(path): + # Include top-level directory in output + info = self.info(path) + sizes[info["name"]] = info["size"] + for f in self.find(path, maxdepth=maxdepth, withdirs=withdirs, **kwargs): + info = self.info(f) + sizes[info["name"]] = info["size"] + if total: + return sum(sizes.values()) + else: + return sizes + + def glob(self, path, maxdepth=None, **kwargs): + """ + Find files by glob-matching. + + If the path ends with '/', only folders are returned. + + We support ``"**"``, + ``"?"`` and ``"[..]"``. We do not support ^ for pattern negation. + + The `maxdepth` option is applied on the first `**` found in the path. + + kwargs are passed to ``ls``. + """ + if maxdepth is not None and maxdepth < 1: + raise ValueError("maxdepth must be at least 1") + + import re + + seps = (os.path.sep, os.path.altsep) if os.path.altsep else (os.path.sep,) + ends_with_sep = path.endswith(seps) # _strip_protocol strips trailing slash + path = self._strip_protocol(path) + append_slash_to_dirname = ends_with_sep or path.endswith( + tuple(sep + "**" for sep in seps) + ) + idx_star = path.find("*") if path.find("*") >= 0 else len(path) + idx_qmark = path.find("?") if path.find("?") >= 0 else len(path) + idx_brace = path.find("[") if path.find("[") >= 0 else len(path) + + min_idx = min(idx_star, idx_qmark, idx_brace) + + detail = kwargs.pop("detail", False) + + if not has_magic(path): + if self.exists(path, **kwargs): + if not detail: + return [path] + else: + return {path: self.info(path, **kwargs)} + else: + if not detail: + return [] # glob of non-existent returns empty + else: + return {} + elif "/" in path[:min_idx]: + min_idx = path[:min_idx].rindex("/") + root = path[: min_idx + 1] + depth = path[min_idx + 1 :].count("/") + 1 + else: + root = "" + depth = path[min_idx + 1 :].count("/") + 1 + + if "**" in path: + if maxdepth is not None: + idx_double_stars = path.find("**") + depth_double_stars = path[idx_double_stars:].count("/") + 1 + depth = depth - depth_double_stars + maxdepth + else: + depth = None + + allpaths = self.find(root, maxdepth=depth, withdirs=True, detail=True, **kwargs) + + pattern = glob_translate(path + ("/" if ends_with_sep else "")) + pattern = re.compile(pattern) + + out = { + p: info + for p, info in sorted(allpaths.items()) + if pattern.match( + ( + p + "/" + if append_slash_to_dirname and info["type"] == "directory" + else p + ) + ) + } + + if detail: + return out + else: + return list(out) + + def exists(self, path, **kwargs): + """Is there a file at the given path""" + try: + self.info(path, **kwargs) + return True + except: # noqa: E722 + # any exception allowed bar FileNotFoundError? + return False + + def lexists(self, path, **kwargs): + """If there is a file at the given path (including + broken links)""" + return self.exists(path) + + def info(self, path, **kwargs): + """Give details of entry at path + + Returns a single dictionary, with exactly the same information as ``ls`` + would with ``detail=True``. + + The default implementation should calls ls and could be overridden by a + shortcut. kwargs are passed on to ```ls()``. + + Some file systems might not be able to measure the file's size, in + which case, the returned dict will include ``'size': None``. + + Returns + ------- + dict with keys: name (full path in the FS), size (in bytes), type (file, + directory, or something else) and other FS-specific keys. + """ + path = self._strip_protocol(path) + out = self.ls(self._parent(path), detail=True, **kwargs) + out = [o for o in out if o["name"].rstrip("/") == path] + if out: + return out[0] + out = self.ls(path, detail=True, **kwargs) + path = path.rstrip("/") + out1 = [o for o in out if o["name"].rstrip("/") == path] + if len(out1) == 1: + if "size" not in out1[0]: + out1[0]["size"] = None + return out1[0] + elif len(out1) > 1 or out: + return {"name": path, "size": 0, "type": "directory"} + else: + raise FileNotFoundError(path) + + def checksum(self, path): + """Unique value for current version of file + + If the checksum is the same from one moment to another, the contents + are guaranteed to be the same. If the checksum changes, the contents + *might* have changed. + + This should normally be overridden; default will probably capture + creation/modification timestamp (which would be good) or maybe + access timestamp (which would be bad) + """ + return int(tokenize(self.info(path)), 16) + + def size(self, path): + """Size in bytes of file""" + return self.info(path).get("size", None) + + def sizes(self, paths): + """Size in bytes of each file in a list of paths""" + return [self.size(p) for p in paths] + + def isdir(self, path): + """Is this entry directory-like?""" + try: + return self.info(path)["type"] == "directory" + except OSError: + return False + + def isfile(self, path): + """Is this entry file-like?""" + try: + return self.info(path)["type"] == "file" + except: # noqa: E722 + return False + + def read_text(self, path, encoding=None, errors=None, newline=None, **kwargs): + """Get the contents of the file as a string. + + Parameters + ---------- + path: str + URL of file on this filesystems + encoding, errors, newline: same as `open`. + """ + with self.open( + path, + mode="r", + encoding=encoding, + errors=errors, + newline=newline, + **kwargs, + ) as f: + return f.read() + + def write_text( + self, path, value, encoding=None, errors=None, newline=None, **kwargs + ): + """Write the text to the given file. + + An existing file will be overwritten. + + Parameters + ---------- + path: str + URL of file on this filesystems + value: str + Text to write. + encoding, errors, newline: same as `open`. + """ + with self.open( + path, + mode="w", + encoding=encoding, + errors=errors, + newline=newline, + **kwargs, + ) as f: + return f.write(value) + + def cat_file(self, path, start=None, end=None, **kwargs): + """Get the content of a file + + Parameters + ---------- + path: URL of file on this filesystems + start, end: int + Bytes limits of the read. If negative, backwards from end, + like usual python slices. Either can be None for start or + end of file, respectively + kwargs: passed to ``open()``. + """ + # explicitly set buffering off? + with self.open(path, "rb", **kwargs) as f: + if start is not None: + if start >= 0: + f.seek(start) + else: + f.seek(max(0, f.size + start)) + if end is not None: + if end < 0: + end = f.size + end + return f.read(end - f.tell()) + return f.read() + + def pipe_file(self, path, value, **kwargs): + """Set the bytes of given file""" + with self.open(path, "wb", **kwargs) as f: + f.write(value) + + def pipe(self, path, value=None, **kwargs): + """Put value into path + + (counterpart to ``cat``) + + Parameters + ---------- + path: string or dict(str, bytes) + If a string, a single remote location to put ``value`` bytes; if a dict, + a mapping of {path: bytesvalue}. + value: bytes, optional + If using a single path, these are the bytes to put there. Ignored if + ``path`` is a dict + """ + if isinstance(path, str): + self.pipe_file(self._strip_protocol(path), value, **kwargs) + elif isinstance(path, dict): + for k, v in path.items(): + self.pipe_file(self._strip_protocol(k), v, **kwargs) + else: + raise ValueError("path must be str or dict") + + def cat_ranges( + self, paths, starts, ends, max_gap=None, on_error="return", **kwargs + ): + """Get the contents of byte ranges from one or more files + + Parameters + ---------- + paths: list + A list of of filepaths on this filesystems + starts, ends: int or list + Bytes limits of the read. If using a single int, the same value will be + used to read all the specified files. + """ + if max_gap is not None: + raise NotImplementedError + if not isinstance(paths, list): + raise TypeError + if not isinstance(starts, list): + starts = [starts] * len(paths) + if not isinstance(ends, list): + ends = [ends] * len(paths) + if len(starts) != len(paths) or len(ends) != len(paths): + raise ValueError + out = [] + for p, s, e in zip(paths, starts, ends): + try: + out.append(self.cat_file(p, s, e)) + except Exception as e: + if on_error == "return": + out.append(e) + else: + raise + return out + + def cat(self, path, recursive=False, on_error="raise", **kwargs): + """Fetch (potentially multiple) paths' contents + + Parameters + ---------- + recursive: bool + If True, assume the path(s) are directories, and get all the + contained files + on_error : "raise", "omit", "return" + If raise, an underlying exception will be raised (converted to KeyError + if the type is in self.missing_exceptions); if omit, keys with exception + will simply not be included in the output; if "return", all keys are + included in the output, but the value will be bytes or an exception + instance. + kwargs: passed to cat_file + + Returns + ------- + dict of {path: contents} if there are multiple paths + or the path has been otherwise expanded + """ + paths = self.expand_path(path, recursive=recursive) + if ( + len(paths) > 1 + or isinstance(path, list) + or paths[0] != self._strip_protocol(path) + ): + out = {} + for path in paths: + try: + out[path] = self.cat_file(path, **kwargs) + except Exception as e: + if on_error == "raise": + raise + if on_error == "return": + out[path] = e + return out + else: + return self.cat_file(paths[0], **kwargs) + + def get_file(self, rpath, lpath, callback=DEFAULT_CALLBACK, outfile=None, **kwargs): + """Copy single remote file to local""" + from .implementations.local import LocalFileSystem + + if isfilelike(lpath): + outfile = lpath + elif self.isdir(rpath): + os.makedirs(lpath, exist_ok=True) + return None + + fs = LocalFileSystem(auto_mkdir=True) + fs.makedirs(fs._parent(lpath), exist_ok=True) + + with self.open(rpath, "rb", **kwargs) as f1: + if outfile is None: + outfile = open(lpath, "wb") + + try: + callback.set_size(getattr(f1, "size", None)) + data = True + while data: + data = f1.read(self.blocksize) + segment_len = outfile.write(data) + if segment_len is None: + segment_len = len(data) + callback.relative_update(segment_len) + finally: + if not isfilelike(lpath): + outfile.close() + + def get( + self, + rpath, + lpath, + recursive=False, + callback=DEFAULT_CALLBACK, + maxdepth=None, + **kwargs, + ): + """Copy file(s) to local. + + Copies a specific file or tree of files (if recursive=True). If lpath + ends with a "/", it will be assumed to be a directory, and target files + will go within. Can submit a list of paths, which may be glob-patterns + and will be expanded. + + Calls get_file for each source. + """ + if isinstance(lpath, list) and isinstance(rpath, list): + # No need to expand paths when both source and destination + # are provided as lists + rpaths = rpath + lpaths = lpath + else: + from .implementations.local import ( + LocalFileSystem, + make_path_posix, + trailing_sep, + ) + + source_is_str = isinstance(rpath, str) + rpaths = self.expand_path(rpath, recursive=recursive, maxdepth=maxdepth) + if source_is_str and (not recursive or maxdepth is not None): + # Non-recursive glob does not copy directories + rpaths = [p for p in rpaths if not (trailing_sep(p) or self.isdir(p))] + if not rpaths: + return + + if isinstance(lpath, str): + lpath = make_path_posix(lpath) + + source_is_file = len(rpaths) == 1 + dest_is_dir = isinstance(lpath, str) and ( + trailing_sep(lpath) or LocalFileSystem().isdir(lpath) + ) + + exists = source_is_str and ( + (has_magic(rpath) and source_is_file) + or (not has_magic(rpath) and dest_is_dir and not trailing_sep(rpath)) + ) + lpaths = other_paths( + rpaths, + lpath, + exists=exists, + flatten=not source_is_str, + ) + + callback.set_size(len(lpaths)) + for lpath, rpath in callback.wrap(zip(lpaths, rpaths)): + with callback.branched(rpath, lpath) as child: + self.get_file(rpath, lpath, callback=child, **kwargs) + + def put_file(self, lpath, rpath, callback=DEFAULT_CALLBACK, **kwargs): + """Copy single file to remote""" + if os.path.isdir(lpath): + self.makedirs(rpath, exist_ok=True) + return None + + with open(lpath, "rb") as f1: + size = f1.seek(0, 2) + callback.set_size(size) + f1.seek(0) + + self.mkdirs(self._parent(os.fspath(rpath)), exist_ok=True) + with self.open(rpath, "wb", **kwargs) as f2: + while f1.tell() < size: + data = f1.read(self.blocksize) + segment_len = f2.write(data) + if segment_len is None: + segment_len = len(data) + callback.relative_update(segment_len) + + def put( + self, + lpath, + rpath, + recursive=False, + callback=DEFAULT_CALLBACK, + maxdepth=None, + **kwargs, + ): + """Copy file(s) from local. + + Copies a specific file or tree of files (if recursive=True). If rpath + ends with a "/", it will be assumed to be a directory, and target files + will go within. + + Calls put_file for each source. + """ + if isinstance(lpath, list) and isinstance(rpath, list): + # No need to expand paths when both source and destination + # are provided as lists + rpaths = rpath + lpaths = lpath + else: + from .implementations.local import ( + LocalFileSystem, + make_path_posix, + trailing_sep, + ) + + source_is_str = isinstance(lpath, str) + if source_is_str: + lpath = make_path_posix(lpath) + fs = LocalFileSystem() + lpaths = fs.expand_path(lpath, recursive=recursive, maxdepth=maxdepth) + if source_is_str and (not recursive or maxdepth is not None): + # Non-recursive glob does not copy directories + lpaths = [p for p in lpaths if not (trailing_sep(p) or fs.isdir(p))] + if not lpaths: + return + + source_is_file = len(lpaths) == 1 + dest_is_dir = isinstance(rpath, str) and ( + trailing_sep(rpath) or self.isdir(rpath) + ) + + rpath = ( + self._strip_protocol(rpath) + if isinstance(rpath, str) + else [self._strip_protocol(p) for p in rpath] + ) + exists = source_is_str and ( + (has_magic(lpath) and source_is_file) + or (not has_magic(lpath) and dest_is_dir and not trailing_sep(lpath)) + ) + rpaths = other_paths( + lpaths, + rpath, + exists=exists, + flatten=not source_is_str, + ) + + callback.set_size(len(rpaths)) + for lpath, rpath in callback.wrap(zip(lpaths, rpaths)): + with callback.branched(lpath, rpath) as child: + self.put_file(lpath, rpath, callback=child, **kwargs) + + def head(self, path, size=1024): + """Get the first ``size`` bytes from file""" + with self.open(path, "rb") as f: + return f.read(size) + + def tail(self, path, size=1024): + """Get the last ``size`` bytes from file""" + with self.open(path, "rb") as f: + f.seek(max(-size, -f.size), 2) + return f.read() + + def cp_file(self, path1, path2, **kwargs): + raise NotImplementedError + + def copy( + self, path1, path2, recursive=False, maxdepth=None, on_error=None, **kwargs + ): + """Copy within two locations in the filesystem + + on_error : "raise", "ignore" + If raise, any not-found exceptions will be raised; if ignore any + not-found exceptions will cause the path to be skipped; defaults to + raise unless recursive is true, where the default is ignore + """ + if on_error is None and recursive: + on_error = "ignore" + elif on_error is None: + on_error = "raise" + + if isinstance(path1, list) and isinstance(path2, list): + # No need to expand paths when both source and destination + # are provided as lists + paths1 = path1 + paths2 = path2 + else: + from .implementations.local import trailing_sep + + source_is_str = isinstance(path1, str) + paths1 = self.expand_path(path1, recursive=recursive, maxdepth=maxdepth) + if source_is_str and (not recursive or maxdepth is not None): + # Non-recursive glob does not copy directories + paths1 = [p for p in paths1 if not (trailing_sep(p) or self.isdir(p))] + if not paths1: + return + + source_is_file = len(paths1) == 1 + dest_is_dir = isinstance(path2, str) and ( + trailing_sep(path2) or self.isdir(path2) + ) + + exists = source_is_str and ( + (has_magic(path1) and source_is_file) + or (not has_magic(path1) and dest_is_dir and not trailing_sep(path1)) + ) + paths2 = other_paths( + paths1, + path2, + exists=exists, + flatten=not source_is_str, + ) + + for p1, p2 in zip(paths1, paths2): + try: + self.cp_file(p1, p2, **kwargs) + except FileNotFoundError: + if on_error == "raise": + raise + + def expand_path(self, path, recursive=False, maxdepth=None, **kwargs): + """Turn one or more globs or directories into a list of all matching paths + to files or directories. + + kwargs are passed to ``glob`` or ``find``, which may in turn call ``ls`` + """ + + if maxdepth is not None and maxdepth < 1: + raise ValueError("maxdepth must be at least 1") + + if isinstance(path, (str, os.PathLike)): + out = self.expand_path([path], recursive, maxdepth) + else: + out = set() + path = [self._strip_protocol(p) for p in path] + for p in path: + if has_magic(p): + bit = set(self.glob(p, maxdepth=maxdepth, **kwargs)) + out |= bit + if recursive: + # glob call above expanded one depth so if maxdepth is defined + # then decrement it in expand_path call below. If it is zero + # after decrementing then avoid expand_path call. + if maxdepth is not None and maxdepth <= 1: + continue + out |= set( + self.expand_path( + list(bit), + recursive=recursive, + maxdepth=maxdepth - 1 if maxdepth is not None else None, + **kwargs, + ) + ) + continue + elif recursive: + rec = set( + self.find( + p, maxdepth=maxdepth, withdirs=True, detail=False, **kwargs + ) + ) + out |= rec + if p not in out and (recursive is False or self.exists(p)): + # should only check once, for the root + out.add(p) + if not out: + raise FileNotFoundError(path) + return sorted(out) + + def mv(self, path1, path2, recursive=False, maxdepth=None, **kwargs): + """Move file(s) from one location to another""" + if path1 == path2: + logger.debug("%s mv: The paths are the same, so no files were moved.", self) + else: + # explicitly raise exception to prevent data corruption + self.copy( + path1, path2, recursive=recursive, maxdepth=maxdepth, onerror="raise" + ) + self.rm(path1, recursive=recursive) + + def rm_file(self, path): + """Delete a file""" + self._rm(path) + + def _rm(self, path): + """Delete one file""" + # this is the old name for the method, prefer rm_file + raise NotImplementedError + + def rm(self, path, recursive=False, maxdepth=None): + """Delete files. + + Parameters + ---------- + path: str or list of str + File(s) to delete. + recursive: bool + If file(s) are directories, recursively delete contents and then + also remove the directory + maxdepth: int or None + Depth to pass to walk for finding files to delete, if recursive. + If None, there will be no limit and infinite recursion may be + possible. + """ + path = self.expand_path(path, recursive=recursive, maxdepth=maxdepth) + for p in reversed(path): + self.rm_file(p) + + @classmethod + def _parent(cls, path): + path = cls._strip_protocol(path) + if "/" in path: + parent = path.rsplit("/", 1)[0].lstrip(cls.root_marker) + return cls.root_marker + parent + else: + return cls.root_marker + + def _open( + self, + path, + mode="rb", + block_size=None, + autocommit=True, + cache_options=None, + **kwargs, + ): + """Return raw bytes-mode file-like from the file-system""" + return AbstractBufferedFile( + self, + path, + mode, + block_size, + autocommit, + cache_options=cache_options, + **kwargs, + ) + + def open( + self, + path, + mode="rb", + block_size=None, + cache_options=None, + compression=None, + **kwargs, + ): + """ + Return a file-like object from the filesystem + + The resultant instance must function correctly in a context ``with`` + block. + + Parameters + ---------- + path: str + Target file + mode: str like 'rb', 'w' + See builtin ``open()`` + block_size: int + Some indication of buffering - this is a value in bytes + cache_options : dict, optional + Extra arguments to pass through to the cache. + compression: string or None + If given, open file using compression codec. Can either be a compression + name (a key in ``fsspec.compression.compr``) or "infer" to guess the + compression from the filename suffix. + encoding, errors, newline: passed on to TextIOWrapper for text mode + """ + import io + + path = self._strip_protocol(path) + if "b" not in mode: + mode = mode.replace("t", "") + "b" + + text_kwargs = { + k: kwargs.pop(k) + for k in ["encoding", "errors", "newline"] + if k in kwargs + } + return io.TextIOWrapper( + self.open( + path, + mode, + block_size=block_size, + cache_options=cache_options, + compression=compression, + **kwargs, + ), + **text_kwargs, + ) + else: + ac = kwargs.pop("autocommit", not self._intrans) + f = self._open( + path, + mode=mode, + block_size=block_size, + autocommit=ac, + cache_options=cache_options, + **kwargs, + ) + if compression is not None: + from fsspec.compression import compr + from fsspec.core import get_compression + + compression = get_compression(path, compression) + compress = compr[compression] + f = compress(f, mode=mode[0]) + + if not ac and "r" not in mode: + self.transaction.files.append(f) + return f + + def touch(self, path, truncate=True, **kwargs): + """Create empty file, or update timestamp + + Parameters + ---------- + path: str + file location + truncate: bool + If True, always set file size to 0; if False, update timestamp and + leave file unchanged, if backend allows this + """ + if truncate or not self.exists(path): + with self.open(path, "wb", **kwargs): + pass + else: + raise NotImplementedError # update timestamp, if possible + + def ukey(self, path): + """Hash of file properties, to tell if it has changed""" + return sha256(str(self.info(path)).encode()).hexdigest() + + def read_block(self, fn, offset, length, delimiter=None): + """Read a block of bytes from + + Starting at ``offset`` of the file, read ``length`` bytes. If + ``delimiter`` is set then we ensure that the read starts and stops at + delimiter boundaries that follow the locations ``offset`` and ``offset + + length``. If ``offset`` is zero then we start at zero. The + bytestring returned WILL include the end delimiter string. + + If offset+length is beyond the eof, reads to eof. + + Parameters + ---------- + fn: string + Path to filename + offset: int + Byte offset to start read + length: int + Number of bytes to read. If None, read to end. + delimiter: bytes (optional) + Ensure reading starts and stops at delimiter bytestring + + Examples + -------- + >>> fs.read_block('data/file.csv', 0, 13) # doctest: +SKIP + b'Alice, 100\\nBo' + >>> fs.read_block('data/file.csv', 0, 13, delimiter=b'\\n') # doctest: +SKIP + b'Alice, 100\\nBob, 200\\n' + + Use ``length=None`` to read to the end of the file. + >>> fs.read_block('data/file.csv', 0, None, delimiter=b'\\n') # doctest: +SKIP + b'Alice, 100\\nBob, 200\\nCharlie, 300' + + See Also + -------- + :func:`fsspec.utils.read_block` + """ + with self.open(fn, "rb") as f: + size = f.size + if length is None: + length = size + if size is not None and offset + length > size: + length = size - offset + return read_block(f, offset, length, delimiter) + + def to_json(self): + """ + JSON representation of this filesystem instance + + Returns + ------- + str: JSON structure with keys cls (the python location of this class), + protocol (text name of this class's protocol, first one in case of + multiple), args (positional args, usually empty), and all other + kwargs as their own keys. + """ + import json + + cls = type(self) + cls = ".".join((cls.__module__, cls.__name__)) + proto = ( + self.protocol[0] + if isinstance(self.protocol, (tuple, list)) + else self.protocol + ) + return json.dumps( + dict( + cls=cls, + protocol=proto, + args=self.storage_args, + **self.storage_options, + ) + ) + + @staticmethod + def from_json(blob): + """ + Recreate a filesystem instance from JSON representation + + See ``.to_json()`` for the expected structure of the input + + Parameters + ---------- + blob: str + + Returns + ------- + file system instance, not necessarily of this particular class. + """ + import json + + from .registry import _import_class, get_filesystem_class + + dic = json.loads(blob) + protocol = dic.pop("protocol") + try: + cls = _import_class(dic.pop("cls")) + except (ImportError, ValueError, RuntimeError, KeyError): + cls = get_filesystem_class(protocol) + return cls(*dic.pop("args", ()), **dic) + + def _get_pyarrow_filesystem(self): + """ + Make a version of the FS instance which will be acceptable to pyarrow + """ + # all instances already also derive from pyarrow + return self + + def get_mapper(self, root="", check=False, create=False, missing_exceptions=None): + """Create key/value store based on this file-system + + Makes a MutableMapping interface to the FS at the given root path. + See ``fsspec.mapping.FSMap`` for further details. + """ + from .mapping import FSMap + + return FSMap( + root, + self, + check=check, + create=create, + missing_exceptions=missing_exceptions, + ) + + @classmethod + def clear_instance_cache(cls): + """ + Clear the cache of filesystem instances. + + Notes + ----- + Unless overridden by setting the ``cachable`` class attribute to False, + the filesystem class stores a reference to newly created instances. This + prevents Python's normal rules around garbage collection from working, + since the instances refcount will not drop to zero until + ``clear_instance_cache`` is called. + """ + cls._cache.clear() + + def created(self, path): + """Return the created timestamp of a file as a datetime.datetime""" + raise NotImplementedError + + def modified(self, path): + """Return the modified timestamp of a file as a datetime.datetime""" + raise NotImplementedError + + # ------------------------------------------------------------------------ + # Aliases + + def read_bytes(self, path, start=None, end=None, **kwargs): + """Alias of `AbstractFileSystem.cat_file`.""" + return self.cat_file(path, start=start, end=end, **kwargs) + + def write_bytes(self, path, value, **kwargs): + """Alias of `AbstractFileSystem.pipe_file`.""" + self.pipe_file(path, value, **kwargs) + + def makedir(self, path, create_parents=True, **kwargs): + """Alias of `AbstractFileSystem.mkdir`.""" + return self.mkdir(path, create_parents=create_parents, **kwargs) + + def mkdirs(self, path, exist_ok=False): + """Alias of `AbstractFileSystem.makedirs`.""" + return self.makedirs(path, exist_ok=exist_ok) + + def listdir(self, path, detail=True, **kwargs): + """Alias of `AbstractFileSystem.ls`.""" + return self.ls(path, detail=detail, **kwargs) + + def cp(self, path1, path2, **kwargs): + """Alias of `AbstractFileSystem.copy`.""" + return self.copy(path1, path2, **kwargs) + + def move(self, path1, path2, **kwargs): + """Alias of `AbstractFileSystem.mv`.""" + return self.mv(path1, path2, **kwargs) + + def stat(self, path, **kwargs): + """Alias of `AbstractFileSystem.info`.""" + return self.info(path, **kwargs) + + def disk_usage(self, path, total=True, maxdepth=None, **kwargs): + """Alias of `AbstractFileSystem.du`.""" + return self.du(path, total=total, maxdepth=maxdepth, **kwargs) + + def rename(self, path1, path2, **kwargs): + """Alias of `AbstractFileSystem.mv`.""" + return self.mv(path1, path2, **kwargs) + + def delete(self, path, recursive=False, maxdepth=None): + """Alias of `AbstractFileSystem.rm`.""" + return self.rm(path, recursive=recursive, maxdepth=maxdepth) + + def upload(self, lpath, rpath, recursive=False, **kwargs): + """Alias of `AbstractFileSystem.put`.""" + return self.put(lpath, rpath, recursive=recursive, **kwargs) + + def download(self, rpath, lpath, recursive=False, **kwargs): + """Alias of `AbstractFileSystem.get`.""" + return self.get(rpath, lpath, recursive=recursive, **kwargs) + + def sign(self, path, expiration=100, **kwargs): + """Create a signed URL representing the given path + + Some implementations allow temporary URLs to be generated, as a + way of delegating credentials. + + Parameters + ---------- + path : str + The path on the filesystem + expiration : int + Number of seconds to enable the URL for (if supported) + + Returns + ------- + URL : str + The signed URL + + Raises + ------ + NotImplementedError : if method is not implemented for a filesystem + """ + raise NotImplementedError("Sign is not implemented for this filesystem") + + def _isfilestore(self): + # Originally inherited from pyarrow DaskFileSystem. Keeping this + # here for backwards compatibility as long as pyarrow uses its + # legacy fsspec-compatible filesystems and thus accepts fsspec + # filesystems as well + return False + + +class AbstractBufferedFile(io.IOBase): + """Convenient class to derive from to provide buffering + + In the case that the backend does not provide a pythonic file-like object + already, this class contains much of the logic to build one. The only + methods that need to be overridden are ``_upload_chunk``, + ``_initiate_upload`` and ``_fetch_range``. + """ + + DEFAULT_BLOCK_SIZE = 5 * 2**20 + _details = None + + def __init__( + self, + fs, + path, + mode="rb", + block_size="default", + autocommit=True, + cache_type="readahead", + cache_options=None, + size=None, + **kwargs, + ): + """ + Template for files with buffered reading and writing + + Parameters + ---------- + fs: instance of FileSystem + path: str + location in file-system + mode: str + Normal file modes. Currently only 'wb', 'ab' or 'rb'. Some file + systems may be read-only, and some may not support append. + block_size: int + Buffer size for reading or writing, 'default' for class default + autocommit: bool + Whether to write to final destination; may only impact what + happens when file is being closed. + cache_type: {"readahead", "none", "mmap", "bytes"}, default "readahead" + Caching policy in read mode. See the definitions in ``core``. + cache_options : dict + Additional options passed to the constructor for the cache specified + by `cache_type`. + size: int + If given and in read mode, suppressed having to look up the file size + kwargs: + Gets stored as self.kwargs + """ + from .core import caches + + self.path = path + self.fs = fs + self.mode = mode + self.blocksize = ( + self.DEFAULT_BLOCK_SIZE if block_size in ["default", None] else block_size + ) + self.loc = 0 + self.autocommit = autocommit + self.end = None + self.start = None + self.closed = False + + if cache_options is None: + cache_options = {} + + if "trim" in kwargs: + warnings.warn( + "Passing 'trim' to control the cache behavior has been deprecated. " + "Specify it within the 'cache_options' argument instead.", + FutureWarning, + ) + cache_options["trim"] = kwargs.pop("trim") + + self.kwargs = kwargs + + if mode not in {"ab", "rb", "wb"}: + raise NotImplementedError("File mode not supported") + if mode == "rb": + if size is not None: + self.size = size + else: + self.size = self.details["size"] + self.cache = caches[cache_type]( + self.blocksize, self._fetch_range, self.size, **cache_options + ) + else: + self.buffer = io.BytesIO() + self.offset = None + self.forced = False + self.location = None + + @property + def details(self): + if self._details is None: + self._details = self.fs.info(self.path) + return self._details + + @details.setter + def details(self, value): + self._details = value + self.size = value["size"] + + @property + def full_name(self): + return _unstrip_protocol(self.path, self.fs) + + @property + def closed(self): + # get around this attr being read-only in IOBase + # use getattr here, since this can be called during del + return getattr(self, "_closed", True) + + @closed.setter + def closed(self, c): + self._closed = c + + def __hash__(self): + if "w" in self.mode: + return id(self) + else: + return int(tokenize(self.details), 16) + + def __eq__(self, other): + """Files are equal if they have the same checksum, only in read mode""" + if self is other: + return True + return ( + isinstance(other, type(self)) + and self.mode == "rb" + and other.mode == "rb" + and hash(self) == hash(other) + ) + + def commit(self): + """Move from temp to final destination""" + + def discard(self): + """Throw away temporary file""" + + def info(self): + """File information about this path""" + if "r" in self.mode: + return self.details + else: + raise ValueError("Info not available while writing") + + def tell(self): + """Current file location""" + return self.loc + + def seek(self, loc, whence=0): + """Set current file location + + Parameters + ---------- + loc: int + byte location + whence: {0, 1, 2} + from start of file, current location or end of file, resp. + """ + loc = int(loc) + if not self.mode == "rb": + raise OSError(ESPIPE, "Seek only available in read mode") + if whence == 0: + nloc = loc + elif whence == 1: + nloc = self.loc + loc + elif whence == 2: + nloc = self.size + loc + else: + raise ValueError(f"invalid whence ({whence}, should be 0, 1 or 2)") + if nloc < 0: + raise ValueError("Seek before start of file") + self.loc = nloc + return self.loc + + def write(self, data): + """ + Write data to buffer. + + Buffer only sent on flush() or if buffer is greater than + or equal to blocksize. + + Parameters + ---------- + data: bytes + Set of bytes to be written. + """ + if self.mode not in {"wb", "ab"}: + raise ValueError("File not in write mode") + if self.closed: + raise ValueError("I/O operation on closed file.") + if self.forced: + raise ValueError("This file has been force-flushed, can only close") + out = self.buffer.write(data) + self.loc += out + if self.buffer.tell() >= self.blocksize: + self.flush() + return out + + def flush(self, force=False): + """ + Write buffered data to backend store. + + Writes the current buffer, if it is larger than the block-size, or if + the file is being closed. + + Parameters + ---------- + force: bool + When closing, write the last block even if it is smaller than + blocks are allowed to be. Disallows further writing to this file. + """ + + if self.closed: + raise ValueError("Flush on closed file") + if force and self.forced: + raise ValueError("Force flush cannot be called more than once") + if force: + self.forced = True + + if self.mode not in {"wb", "ab"}: + # no-op to flush on read-mode + return + + if not force and self.buffer.tell() < self.blocksize: + # Defer write on small block + return + + if self.offset is None: + # Initialize a multipart upload + self.offset = 0 + try: + self._initiate_upload() + except: # noqa: E722 + self.closed = True + raise + + if self._upload_chunk(final=force) is not False: + self.offset += self.buffer.seek(0, 2) + self.buffer = io.BytesIO() + + def _upload_chunk(self, final=False): + """Write one part of a multi-block file upload + + Parameters + ========== + final: bool + This is the last block, so should complete file, if + self.autocommit is True. + """ + # may not yet have been initialized, may need to call _initialize_upload + + def _initiate_upload(self): + """Create remote file/upload""" + pass + + def _fetch_range(self, start, end): + """Get the specified set of bytes from remote""" + raise NotImplementedError + + def read(self, length=-1): + """ + Return data from cache, or fetch pieces as necessary + + Parameters + ---------- + length: int (-1) + Number of bytes to read; if <0, all remaining bytes. + """ + length = -1 if length is None else int(length) + if self.mode != "rb": + raise ValueError("File not in read mode") + if length < 0: + length = self.size - self.loc + if self.closed: + raise ValueError("I/O operation on closed file.") + if length == 0: + # don't even bother calling fetch + return b"" + out = self.cache._fetch(self.loc, self.loc + length) + + logger.debug( + "%s read: %i - %i %s", + self, + self.loc, + self.loc + length, + self.cache._log_stats(), + ) + self.loc += len(out) + return out + + def readinto(self, b): + """mirrors builtin file's readinto method + + https://docs.python.org/3/library/io.html#io.RawIOBase.readinto + """ + out = memoryview(b).cast("B") + data = self.read(out.nbytes) + out[: len(data)] = data + return len(data) + + def readuntil(self, char=b"\n", blocks=None): + """Return data between current position and first occurrence of char + + char is included in the output, except if the end of the tile is + encountered first. + + Parameters + ---------- + char: bytes + Thing to find + blocks: None or int + How much to read in each go. Defaults to file blocksize - which may + mean a new read on every call. + """ + out = [] + while True: + start = self.tell() + part = self.read(blocks or self.blocksize) + if len(part) == 0: + break + found = part.find(char) + if found > -1: + out.append(part[: found + len(char)]) + self.seek(start + found + len(char)) + break + out.append(part) + return b"".join(out) + + def readline(self): + """Read until first occurrence of newline character + + Note that, because of character encoding, this is not necessarily a + true line ending. + """ + return self.readuntil(b"\n") + + def __next__(self): + out = self.readline() + if out: + return out + raise StopIteration + + def __iter__(self): + return self + + def readlines(self): + """Return all data, split by the newline character""" + data = self.read() + lines = data.split(b"\n") + out = [l + b"\n" for l in lines[:-1]] + if data.endswith(b"\n"): + return out + else: + return out + [lines[-1]] + # return list(self) ??? + + def readinto1(self, b): + return self.readinto(b) + + def close(self): + """Close file + + Finalizes writes, discards cache + """ + if getattr(self, "_unclosable", False): + return + if self.closed: + return + if self.mode == "rb": + self.cache = None + else: + if not self.forced: + self.flush(force=True) + + if self.fs is not None: + self.fs.invalidate_cache(self.path) + self.fs.invalidate_cache(self.fs._parent(self.path)) + + self.closed = True + + def readable(self): + """Whether opened for reading""" + return self.mode == "rb" and not self.closed + + def seekable(self): + """Whether is seekable (only in read mode)""" + return self.readable() + + def writable(self): + """Whether opened for writing""" + return self.mode in {"wb", "ab"} and not self.closed + + def __del__(self): + if not self.closed: + self.close() + + def __str__(self): + return f"" + + __repr__ = __str__ + + def __enter__(self): + return self + + def __exit__(self, *args): + self.close() diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/hf_xet-1.3.2.dist-info/INSTALLER b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/hf_xet-1.3.2.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/hf_xet-1.3.2.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/hf_xet-1.3.2.dist-info/METADATA b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/hf_xet-1.3.2.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..eed417e382589ca226f658db5da489d93802d3a4 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/hf_xet-1.3.2.dist-info/METADATA @@ -0,0 +1,87 @@ +Metadata-Version: 2.4 +Name: hf-xet +Version: 1.3.2 +Classifier: Development Status :: 5 - Production/Stable +Classifier: License :: OSI Approved :: Apache Software License +Classifier: Programming Language :: Rust +Classifier: Programming Language :: Python :: Implementation :: CPython +Classifier: Programming Language :: Python :: Implementation :: PyPy +Classifier: Programming Language :: Python :: 3 +Classifier: Programming Language :: Python :: 3 :: Only +Classifier: Programming Language :: Python :: 3.8 +Classifier: Programming Language :: Python :: 3.9 +Classifier: Programming Language :: Python :: 3.10 +Classifier: Programming Language :: Python :: 3.11 +Classifier: Programming Language :: Python :: 3.12 +Classifier: Programming Language :: Python :: 3.13 +Classifier: Programming Language :: Python :: 3.14 +Classifier: Programming Language :: Python :: Free Threading +Classifier: Programming Language :: Python :: Free Threading :: 2 - Beta +Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence +Requires-Dist: pytest ; extra == 'tests' +Provides-Extra: tests +License-File: LICENSE +Summary: Fast transfer of large files with the Hugging Face Hub. +Maintainer-email: Rajat Arya , Jared Sulzdorf , Di Xiao , Assaf Vayner , Hoyt Koepke +License-Expression: Apache-2.0 +Requires-Python: >=3.8 +Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM +Project-URL: Documentation, https://huggingface.co/docs/hub/xet/index +Project-URL: Homepage, https://github.com/huggingface/xet-core +Project-URL: Issues, https://github.com/huggingface/xet-core/issues +Project-URL: Repository, https://github.com/huggingface/xet-core.git + + +

+ License + GitHub release + Contributor Covenant +

+ +

+

🤗 hf-xet - xet client tech, used in huggingface_hub

+

+ +## Welcome + +`hf-xet` enables `huggingface_hub` to utilize xet storage for uploading and downloading to HF Hub. Xet storage provides chunk-based deduplication, efficient storage/retrieval with local disk caching, and backwards compatibility with Git LFS. This library is not meant to be used directly, and is instead intended to be used from [huggingface_hub](https://pypi.org/project/huggingface-hub). + +## Key features + +♻ **chunk-based deduplication implementation**: avoid transferring and storing chunks that are shared across binary files (models, datasets, etc). + +🤗 **Python bindings**: bindings for [huggingface_hub](https://github.com/huggingface/huggingface_hub/) package. + +↔ **network communications**: concurrent communication to HF Hub Xet backend services (CAS). + +🔖 **local disk caching**: chunk-based cache that sits alongside the existing [huggingface_hub disk cache](https://huggingface.co/docs/huggingface_hub/guides/manage-cache). + +## Installation + +Install the `hf_xet` package with [pip](https://pypi.org/project/hf-xet/): + +```bash +pip install hf_xet +``` + +## Quick Start + +`hf_xet` is not intended to be run independently as it is expected to be used from `huggingface_hub`, so to get started with `huggingface_hub` check out the documentation [here]("https://hf.co/docs/huggingface_hub"). + +## Contributions (feature requests, bugs, etc.) are encouraged & appreciated 💙💚💛💜🧡❤️ + +Please join us in making hf-xet better. We value everyone's contributions. Code is not the only way to help. Answering questions, helping each other, improving documentation, filing issues all help immensely. If you are interested in contributing (please do!), check out the [contribution guide](https://github.com/huggingface/xet-core/blob/main/CONTRIBUTING.md) for this repository. diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/hf_xet-1.3.2.dist-info/RECORD b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/hf_xet-1.3.2.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..3943f0b7838d0a67f488c360d8eaa244896318bb --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/hf_xet-1.3.2.dist-info/RECORD @@ -0,0 +1,9 @@ +hf_xet-1.3.2.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +hf_xet-1.3.2.dist-info/METADATA,sha256=VdBkvkyMsWS6C51XipEai-iZGPIt9ydUqE6iReFQdVI,4882 +hf_xet-1.3.2.dist-info/RECORD,, +hf_xet-1.3.2.dist-info/WHEEL,sha256=ycZRBBQNh-_vBSCR_X0ck7NlL8NDEI7G-tWH-Cjb8uw,143 +hf_xet-1.3.2.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357 +hf_xet-1.3.2.dist-info/sboms/hf_xet.cyclonedx.json,sha256=Z8lgnfjLuQ1OS7SIhJmZNn9qdETBotCqRMS1kgQOnxE,394826 +hf_xet/__init__.py,sha256=E8UDdyQ8glZ_nve9hHEf22bPang8-RKx4VuApXYeQUo,107 +hf_xet/__pycache__/__init__.cpython-312.pyc,, +hf_xet/hf_xet.abi3.so,sha256=35NZSsDbLEwZbmB-1aRzdjHhZD_MwMLDqP-tmCeN2gE,10594520 diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/hf_xet-1.3.2.dist-info/WHEEL b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/hf_xet-1.3.2.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..e3c7eedd8e94fbe7277127c82058423fa8d80b38 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/hf_xet-1.3.2.dist-info/WHEEL @@ -0,0 +1,5 @@ +Wheel-Version: 1.0 +Generator: maturin (1.12.4) +Root-Is-Purelib: false +Tag: cp37-abi3-manylinux_2_17_x86_64 +Tag: cp37-abi3-manylinux2014_x86_64 diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/idna-3.10.dist-info/INSTALLER b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/idna-3.10.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/idna-3.10.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/idna-3.10.dist-info/LICENSE.md b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/idna-3.10.dist-info/LICENSE.md new file mode 100644 index 0000000000000000000000000000000000000000..19b6b45242c16a1025465309eec2ca5009319de3 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/idna-3.10.dist-info/LICENSE.md @@ -0,0 +1,31 @@ +BSD 3-Clause License + +Copyright (c) 2013-2024, Kim Davies and contributors. +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are +met: + +1. Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + +2. 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. + +3. Neither the name of the copyright holder 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. diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/idna-3.10.dist-info/METADATA b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/idna-3.10.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..c42623e9423c23b555d9d352bc5dab518ede02c2 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/idna-3.10.dist-info/METADATA @@ -0,0 +1,250 @@ +Metadata-Version: 2.1 +Name: idna +Version: 3.10 +Summary: Internationalized Domain Names in Applications (IDNA) +Author-email: Kim Davies +Requires-Python: >=3.6 +Description-Content-Type: text/x-rst +Classifier: Development Status :: 5 - Production/Stable +Classifier: Intended Audience :: Developers +Classifier: Intended Audience :: System Administrators +Classifier: License :: OSI Approved :: BSD License +Classifier: Operating System :: OS Independent +Classifier: Programming Language :: Python +Classifier: Programming Language :: Python :: 3 +Classifier: Programming Language :: Python :: 3 :: Only +Classifier: Programming Language :: Python :: 3.6 +Classifier: Programming Language :: Python :: 3.7 +Classifier: Programming Language :: Python :: 3.8 +Classifier: Programming Language :: Python :: 3.9 +Classifier: Programming Language :: Python :: 3.10 +Classifier: Programming Language :: Python :: 3.11 +Classifier: Programming Language :: Python :: 3.12 +Classifier: Programming Language :: Python :: 3.13 +Classifier: Programming Language :: Python :: Implementation :: CPython +Classifier: Programming Language :: Python :: Implementation :: PyPy +Classifier: Topic :: Internet :: Name Service (DNS) +Classifier: Topic :: Software Development :: Libraries :: Python Modules +Classifier: Topic :: Utilities +Requires-Dist: ruff >= 0.6.2 ; extra == "all" +Requires-Dist: mypy >= 1.11.2 ; extra == "all" +Requires-Dist: pytest >= 8.3.2 ; extra == "all" +Requires-Dist: flake8 >= 7.1.1 ; extra == "all" +Project-URL: Changelog, https://github.com/kjd/idna/blob/master/HISTORY.rst +Project-URL: Issue tracker, https://github.com/kjd/idna/issues +Project-URL: Source, https://github.com/kjd/idna +Provides-Extra: all + +Internationalized Domain Names in Applications (IDNA) +===================================================== + +Support for the Internationalized Domain Names in +Applications (IDNA) protocol as specified in `RFC 5891 +`_. This is the latest version of +the protocol and is sometimes referred to as “IDNA 2008”. + +This library also provides support for Unicode Technical +Standard 46, `Unicode IDNA Compatibility Processing +`_. + +This acts as a suitable replacement for the “encodings.idna” +module that comes with the Python standard library, but which +only supports the older superseded IDNA specification (`RFC 3490 +`_). + +Basic functions are simply executed: + +.. code-block:: pycon + + >>> import idna + >>> idna.encode('ドメイン.テスト') + b'xn--eckwd4c7c.xn--zckzah' + >>> print(idna.decode('xn--eckwd4c7c.xn--zckzah')) + ドメイン.テスト + + +Installation +------------ + +This package is available for installation from PyPI: + +.. code-block:: bash + + $ python3 -m pip install idna + + +Usage +----- + +For typical usage, the ``encode`` and ``decode`` functions will take a +domain name argument and perform a conversion to A-labels or U-labels +respectively. + +.. code-block:: pycon + + >>> import idna + >>> idna.encode('ドメイン.テスト') + b'xn--eckwd4c7c.xn--zckzah' + >>> print(idna.decode('xn--eckwd4c7c.xn--zckzah')) + ドメイン.テスト + +You may use the codec encoding and decoding methods using the +``idna.codec`` module: + +.. code-block:: pycon + + >>> import idna.codec + >>> print('домен.испытание'.encode('idna2008')) + b'xn--d1acufc.xn--80akhbyknj4f' + >>> print(b'xn--d1acufc.xn--80akhbyknj4f'.decode('idna2008')) + домен.испытание + +Conversions can be applied at a per-label basis using the ``ulabel`` or +``alabel`` functions if necessary: + +.. code-block:: pycon + + >>> idna.alabel('测试') + b'xn--0zwm56d' + +Compatibility Mapping (UTS #46) ++++++++++++++++++++++++++++++++ + +As described in `RFC 5895 `_, the +IDNA specification does not normalize input from different potential +ways a user may input a domain name. This functionality, known as +a “mapping”, is considered by the specification to be a local +user-interface issue distinct from IDNA conversion functionality. + +This library provides one such mapping that was developed by the +Unicode Consortium. Known as `Unicode IDNA Compatibility Processing +`_, it provides for both a regular +mapping for typical applications, as well as a transitional mapping to +help migrate from older IDNA 2003 applications. Strings are +preprocessed according to Section 4.4 “Preprocessing for IDNA2008” +prior to the IDNA operations. + +For example, “Königsgäßchen” is not a permissible label as *LATIN +CAPITAL LETTER K* is not allowed (nor are capital letters in general). +UTS 46 will convert this into lower case prior to applying the IDNA +conversion. + +.. code-block:: pycon + + >>> import idna + >>> idna.encode('Königsgäßchen') + ... + idna.core.InvalidCodepoint: Codepoint U+004B at position 1 of 'Königsgäßchen' not allowed + >>> idna.encode('Königsgäßchen', uts46=True) + b'xn--knigsgchen-b4a3dun' + >>> print(idna.decode('xn--knigsgchen-b4a3dun')) + königsgäßchen + +Transitional processing provides conversions to help transition from +the older 2003 standard to the current standard. For example, in the +original IDNA specification, the *LATIN SMALL LETTER SHARP S* (ß) was +converted into two *LATIN SMALL LETTER S* (ss), whereas in the current +IDNA specification this conversion is not performed. + +.. code-block:: pycon + + >>> idna.encode('Königsgäßchen', uts46=True, transitional=True) + 'xn--knigsgsschen-lcb0w' + +Implementers should use transitional processing with caution, only in +rare cases where conversion from legacy labels to current labels must be +performed (i.e. IDNA implementations that pre-date 2008). For typical +applications that just need to convert labels, transitional processing +is unlikely to be beneficial and could produce unexpected incompatible +results. + +``encodings.idna`` Compatibility +++++++++++++++++++++++++++++++++ + +Function calls from the Python built-in ``encodings.idna`` module are +mapped to their IDNA 2008 equivalents using the ``idna.compat`` module. +Simply substitute the ``import`` clause in your code to refer to the new +module name. + +Exceptions +---------- + +All errors raised during the conversion following the specification +should raise an exception derived from the ``idna.IDNAError`` base +class. + +More specific exceptions that may be generated as ``idna.IDNABidiError`` +when the error reflects an illegal combination of left-to-right and +right-to-left characters in a label; ``idna.InvalidCodepoint`` when +a specific codepoint is an illegal character in an IDN label (i.e. +INVALID); and ``idna.InvalidCodepointContext`` when the codepoint is +illegal based on its positional context (i.e. it is CONTEXTO or CONTEXTJ +but the contextual requirements are not satisfied.) + +Building and Diagnostics +------------------------ + +The IDNA and UTS 46 functionality relies upon pre-calculated lookup +tables for performance. These tables are derived from computing against +eligibility criteria in the respective standards. These tables are +computed using the command-line script ``tools/idna-data``. + +This tool will fetch relevant codepoint data from the Unicode repository +and perform the required calculations to identify eligibility. There are +three main modes: + +* ``idna-data make-libdata``. Generates ``idnadata.py`` and + ``uts46data.py``, the pre-calculated lookup tables used for IDNA and + UTS 46 conversions. Implementers who wish to track this library against + a different Unicode version may use this tool to manually generate a + different version of the ``idnadata.py`` and ``uts46data.py`` files. + +* ``idna-data make-table``. Generate a table of the IDNA disposition + (e.g. PVALID, CONTEXTJ, CONTEXTO) in the format found in Appendix + B.1 of RFC 5892 and the pre-computed tables published by `IANA + `_. + +* ``idna-data U+0061``. Prints debugging output on the various + properties associated with an individual Unicode codepoint (in this + case, U+0061), that are used to assess the IDNA and UTS 46 status of a + codepoint. This is helpful in debugging or analysis. + +The tool accepts a number of arguments, described using ``idna-data +-h``. Most notably, the ``--version`` argument allows the specification +of the version of Unicode to be used in computing the table data. For +example, ``idna-data --version 9.0.0 make-libdata`` will generate +library data against Unicode 9.0.0. + + +Additional Notes +---------------- + +* **Packages**. The latest tagged release version is published in the + `Python Package Index `_. + +* **Version support**. This library supports Python 3.6 and higher. + As this library serves as a low-level toolkit for a variety of + applications, many of which strive for broad compatibility with older + Python versions, there is no rush to remove older interpreter support. + Removing support for older versions should be well justified in that the + maintenance burden has become too high. + +* **Python 2**. Python 2 is supported by version 2.x of this library. + Use "idna<3" in your requirements file if you need this library for + a Python 2 application. Be advised that these versions are no longer + actively developed. + +* **Testing**. The library has a test suite based on each rule of the + IDNA specification, as well as tests that are provided as part of the + Unicode Technical Standard 46, `Unicode IDNA Compatibility Processing + `_. + +* **Emoji**. It is an occasional request to support emoji domains in + this library. Encoding of symbols like emoji is expressly prohibited by + the technical standard IDNA 2008 and emoji domains are broadly phased + out across the domain industry due to associated security risks. For + now, applications that need to support these non-compliant labels + may wish to consider trying the encode/decode operation in this library + first, and then falling back to using `encodings.idna`. See `the Github + project `_ for more discussion. + diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/idna-3.10.dist-info/RECORD b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/idna-3.10.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..9cfce7f9f4c1a64d85642f865f9d0a3d138a82c6 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/idna-3.10.dist-info/RECORD @@ -0,0 +1,22 @@ +idna-3.10.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +idna-3.10.dist-info/LICENSE.md,sha256=pZ8LDvNjWHQQmkRhykT_enDVBpboFHZ7-vch1Mmw2w8,1541 +idna-3.10.dist-info/METADATA,sha256=URR5ZyDfQ1PCEGhkYoojqfi2Ra0tau2--lhwG4XSfjI,10158 +idna-3.10.dist-info/RECORD,, +idna-3.10.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81 +idna/__init__.py,sha256=MPqNDLZbXqGaNdXxAFhiqFPKEQXju2jNQhCey6-5eJM,868 +idna/__pycache__/__init__.cpython-312.pyc,, +idna/__pycache__/codec.cpython-312.pyc,, +idna/__pycache__/compat.cpython-312.pyc,, +idna/__pycache__/core.cpython-312.pyc,, +idna/__pycache__/idnadata.cpython-312.pyc,, +idna/__pycache__/intranges.cpython-312.pyc,, +idna/__pycache__/package_data.cpython-312.pyc,, +idna/__pycache__/uts46data.cpython-312.pyc,, +idna/codec.py,sha256=PEew3ItwzjW4hymbasnty2N2OXvNcgHB-JjrBuxHPYY,3422 +idna/compat.py,sha256=RzLy6QQCdl9784aFhb2EX9EKGCJjg0P3PilGdeXXcx8,316 +idna/core.py,sha256=YJYyAMnwiQEPjVC4-Fqu_p4CJ6yKKuDGmppBNQNQpFs,13239 +idna/idnadata.py,sha256=W30GcIGvtOWYwAjZj4ZjuouUutC6ffgNuyjJy7fZ-lo,78306 +idna/intranges.py,sha256=amUtkdhYcQG8Zr-CoMM_kVRacxkivC1WgxN1b63KKdU,1898 +idna/package_data.py,sha256=q59S3OXsc5VI8j6vSD0sGBMyk6zZ4vWFREE88yCJYKs,21 +idna/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +idna/uts46data.py,sha256=rt90K9J40gUSwppDPCrhjgi5AA6pWM65dEGRSf6rIhM,239289 diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/idna-3.10.dist-info/WHEEL b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/idna-3.10.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..3b5e64b5e6c4a210201d1676a891fd57b15cda99 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/idna-3.10.dist-info/WHEEL @@ -0,0 +1,4 @@ +Wheel-Version: 1.0 +Generator: flit 3.9.0 +Root-Is-Purelib: true +Tag: py3-none-any diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/nvidia_nccl_cu12-2.27.5.dist-info/INSTALLER b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/nvidia_nccl_cu12-2.27.5.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/nvidia_nccl_cu12-2.27.5.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/nvidia_nccl_cu12-2.27.5.dist-info/METADATA b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/nvidia_nccl_cu12-2.27.5.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..267439298875d4fb658a6913dc51b48c9026e3ea --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/nvidia_nccl_cu12-2.27.5.dist-info/METADATA @@ -0,0 +1,45 @@ +Metadata-Version: 2.4 +Name: nvidia-nccl-cu12 +Version: 2.27.5 +Summary: NVIDIA Collective Communication Library (NCCL) Runtime +Home-page: https://developer.nvidia.com/cuda-zone +Author: Nvidia CUDA Installer Team +Author-email: compute_installer@nvidia.com +License: BSD-3-Clause +Keywords: cuda,nvidia,runtime,machine learning,deep learning +Classifier: Development Status :: 4 - Beta +Classifier: Intended Audience :: Developers +Classifier: Intended Audience :: Education +Classifier: Intended Audience :: Science/Research +Classifier: License :: Other/Proprietary License +Classifier: Natural Language :: English +Classifier: Programming Language :: Python :: 3 +Classifier: Programming Language :: Python :: 3.5 +Classifier: Programming Language :: Python :: 3.6 +Classifier: Programming Language :: Python :: 3.7 +Classifier: Programming Language :: Python :: 3.8 +Classifier: Programming Language :: Python :: 3.9 +Classifier: Programming Language :: Python :: 3.10 +Classifier: Programming Language :: Python :: 3.11 +Classifier: Programming Language :: Python :: 3 :: Only +Classifier: Topic :: Scientific/Engineering +Classifier: Topic :: Scientific/Engineering :: Mathematics +Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence +Classifier: Topic :: Software Development +Classifier: Topic :: Software Development :: Libraries +Classifier: Operating System :: Microsoft :: Windows +Classifier: Operating System :: POSIX :: Linux +Requires-Python: >=3 +License-File: License.txt +Dynamic: author +Dynamic: author-email +Dynamic: classifier +Dynamic: description +Dynamic: home-page +Dynamic: keywords +Dynamic: license +Dynamic: license-file +Dynamic: requires-python +Dynamic: summary + +NCCL (pronounced "Nickel") is a stand-alone library of standard collective communication routines for GPUs, implementing all-reduce, all-gather, reduce, broadcast, and reduce-scatter. It has been optimized to achieve high bandwidth on any platform using PCIe, NVLink, NVswitch, as well as networking using InfiniBand Verbs or TCP/IP sockets. diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/nvidia_nccl_cu12-2.27.5.dist-info/RECORD b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/nvidia_nccl_cu12-2.27.5.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..d5cd7999f23e3c076bf3e6003af961d30e12e933 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/nvidia_nccl_cu12-2.27.5.dist-info/RECORD @@ -0,0 +1,8 @@ +nvidia/nccl/include/nccl.h,sha256=SyQKH6eYWMqM3_z1wvxArqyos-rsYXKCE3FfCQKro6g,22566 +nvidia/nccl/lib/libnccl.so.2,sha256=HZwQMi76zeaVz2QNsfKfwVhskp8Qh_DhIC0tXnD7cWg,429628328 +nvidia_nccl_cu12-2.27.5.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +nvidia_nccl_cu12-2.27.5.dist-info/METADATA,sha256=8QYX3uAOzTd6A-kQPrzyDDMwXQ30cZGlQEmKtF-Kxn0,2019 +nvidia_nccl_cu12-2.27.5.dist-info/RECORD,, +nvidia_nccl_cu12-2.27.5.dist-info/WHEEL,sha256=gy6FWQgpujK_dnYc155G2NL32NQjpi5ebTEXjh8SGZQ,144 +nvidia_nccl_cu12-2.27.5.dist-info/licenses/License.txt,sha256=DwF0prTgszrCY3W_cpUzB1sy9MUaW2gCo9dC19zcmnY,1895 +nvidia_nccl_cu12-2.27.5.dist-info/top_level.txt,sha256=fTkAtiFuL16nUrB9ytDDtpytz2t0B4NvYTnRzwAhO14,7 diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/nvidia_nccl_cu12-2.27.5.dist-info/WHEEL b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/nvidia_nccl_cu12-2.27.5.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..29738220b339f398b2cbbc9ac8ccb38a39bdd7b2 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/nvidia_nccl_cu12-2.27.5.dist-info/WHEEL @@ -0,0 +1,6 @@ +Wheel-Version: 1.0 +Generator: setuptools (80.9.0) +Root-Is-Purelib: true +Tag: py3-none-manylinux2014_x86_64 +Tag: py3-none-manylinux_2_17_x86_64 + diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/nvidia_nccl_cu12-2.27.5.dist-info/top_level.txt b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/nvidia_nccl_cu12-2.27.5.dist-info/top_level.txt new file mode 100644 index 0000000000000000000000000000000000000000..862f7abf232cdfbb928609856247292e81c9decb --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/nvidia_nccl_cu12-2.27.5.dist-info/top_level.txt @@ -0,0 +1 @@ +nvidia diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/psutil/__init__.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/psutil/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..880e1a9ddfc1a95b9388ee4d6c1b06c9d53dd6f1 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/psutil/__init__.py @@ -0,0 +1,2485 @@ +# Copyright (c) 2009, Giampaolo Rodola'. All rights reserved. +# Use of this source code is governed by a BSD-style license that can be +# found in the LICENSE file. + +"""psutil is a cross-platform library for retrieving information on +running processes and system utilization (CPU, memory, disks, network, +sensors) in Python. Supported platforms: + + - Linux + - Windows + - macOS + - FreeBSD + - OpenBSD + - NetBSD + - Sun Solaris + - AIX + +Supported Python versions are cPython 3.6+ and PyPy. +""" + +import collections +import contextlib +import datetime +import functools +import os +import signal +import socket +import subprocess +import sys +import threading +import time + +try: + import pwd +except ImportError: + pwd = None + +from . import _common +from . import _ntuples as _ntp +from ._common import AIX +from ._common import BSD +from ._common import CONN_CLOSE +from ._common import CONN_CLOSE_WAIT +from ._common import CONN_CLOSING +from ._common import CONN_ESTABLISHED +from ._common import CONN_FIN_WAIT1 +from ._common import CONN_FIN_WAIT2 +from ._common import CONN_LAST_ACK +from ._common import CONN_LISTEN +from ._common import CONN_NONE +from ._common import CONN_SYN_RECV +from ._common import CONN_SYN_SENT +from ._common import CONN_TIME_WAIT +from ._common import FREEBSD +from ._common import LINUX +from ._common import MACOS +from ._common import NETBSD +from ._common import NIC_DUPLEX_FULL +from ._common import NIC_DUPLEX_HALF +from ._common import NIC_DUPLEX_UNKNOWN +from ._common import OPENBSD +from ._common import OSX # deprecated alias +from ._common import POSIX +from ._common import POWER_TIME_UNKNOWN +from ._common import POWER_TIME_UNLIMITED +from ._common import STATUS_DEAD +from ._common import STATUS_DISK_SLEEP +from ._common import STATUS_IDLE +from ._common import STATUS_LOCKED +from ._common import STATUS_PARKED +from ._common import STATUS_RUNNING +from ._common import STATUS_SLEEPING +from ._common import STATUS_STOPPED +from ._common import STATUS_TRACING_STOP +from ._common import STATUS_WAITING +from ._common import STATUS_WAKING +from ._common import STATUS_ZOMBIE +from ._common import SUNOS +from ._common import WINDOWS +from ._common import AccessDenied +from ._common import Error +from ._common import NoSuchProcess +from ._common import TimeoutExpired +from ._common import ZombieProcess +from ._common import debug +from ._common import memoize_when_activated +from ._common import wrap_numbers as _wrap_numbers + +if LINUX: + # This is public API and it will be retrieved from _pslinux.py + # via sys.modules. + PROCFS_PATH = "/proc" + + from . import _pslinux as _psplatform + from ._pslinux import IOPRIO_CLASS_BE # noqa: F401 + from ._pslinux import IOPRIO_CLASS_IDLE # noqa: F401 + from ._pslinux import IOPRIO_CLASS_NONE # noqa: F401 + from ._pslinux import IOPRIO_CLASS_RT # noqa: F401 + +elif WINDOWS: + from . import _pswindows as _psplatform + from ._psutil_windows import ABOVE_NORMAL_PRIORITY_CLASS # noqa: F401 + from ._psutil_windows import BELOW_NORMAL_PRIORITY_CLASS # noqa: F401 + from ._psutil_windows import HIGH_PRIORITY_CLASS # noqa: F401 + from ._psutil_windows import IDLE_PRIORITY_CLASS # noqa: F401 + from ._psutil_windows import NORMAL_PRIORITY_CLASS # noqa: F401 + from ._psutil_windows import REALTIME_PRIORITY_CLASS # noqa: F401 + from ._pswindows import CONN_DELETE_TCB # noqa: F401 + from ._pswindows import IOPRIO_HIGH # noqa: F401 + from ._pswindows import IOPRIO_LOW # noqa: F401 + from ._pswindows import IOPRIO_NORMAL # noqa: F401 + from ._pswindows import IOPRIO_VERYLOW # noqa: F401 + +elif MACOS: + from . import _psosx as _psplatform + +elif BSD: + from . import _psbsd as _psplatform + +elif SUNOS: + from . import _pssunos as _psplatform + from ._pssunos import CONN_BOUND # noqa: F401 + from ._pssunos import CONN_IDLE # noqa: F401 + + # This is public writable API which is read from _pslinux.py and + # _pssunos.py via sys.modules. + PROCFS_PATH = "/proc" + +elif AIX: + from . import _psaix as _psplatform + + # This is public API and it will be retrieved from _pslinux.py + # via sys.modules. + PROCFS_PATH = "/proc" + +else: # pragma: no cover + msg = f"platform {sys.platform} is not supported" + raise NotImplementedError(msg) + + +# fmt: off +__all__ = [ + # exceptions + "Error", "NoSuchProcess", "ZombieProcess", "AccessDenied", + "TimeoutExpired", + + # constants + "version_info", "__version__", + + "STATUS_RUNNING", "STATUS_IDLE", "STATUS_SLEEPING", "STATUS_DISK_SLEEP", + "STATUS_STOPPED", "STATUS_TRACING_STOP", "STATUS_ZOMBIE", "STATUS_DEAD", + "STATUS_WAKING", "STATUS_LOCKED", "STATUS_WAITING", "STATUS_LOCKED", + "STATUS_PARKED", + + "CONN_ESTABLISHED", "CONN_SYN_SENT", "CONN_SYN_RECV", "CONN_FIN_WAIT1", + "CONN_FIN_WAIT2", "CONN_TIME_WAIT", "CONN_CLOSE", "CONN_CLOSE_WAIT", + "CONN_LAST_ACK", "CONN_LISTEN", "CONN_CLOSING", "CONN_NONE", + # "CONN_IDLE", "CONN_BOUND", + + "AF_LINK", + + "NIC_DUPLEX_FULL", "NIC_DUPLEX_HALF", "NIC_DUPLEX_UNKNOWN", + + "POWER_TIME_UNKNOWN", "POWER_TIME_UNLIMITED", + + "BSD", "FREEBSD", "LINUX", "NETBSD", "OPENBSD", "MACOS", "OSX", "POSIX", + "SUNOS", "WINDOWS", "AIX", + + # "RLIM_INFINITY", "RLIMIT_AS", "RLIMIT_CORE", "RLIMIT_CPU", "RLIMIT_DATA", + # "RLIMIT_FSIZE", "RLIMIT_LOCKS", "RLIMIT_MEMLOCK", "RLIMIT_NOFILE", + # "RLIMIT_NPROC", "RLIMIT_RSS", "RLIMIT_STACK", "RLIMIT_MSGQUEUE", + # "RLIMIT_NICE", "RLIMIT_RTPRIO", "RLIMIT_RTTIME", "RLIMIT_SIGPENDING", + + # classes + "Process", "Popen", + + # functions + "pid_exists", "pids", "process_iter", "wait_procs", # proc + "virtual_memory", "swap_memory", # memory + "cpu_times", "cpu_percent", "cpu_times_percent", "cpu_count", # cpu + "cpu_stats", # "cpu_freq", "getloadavg" + "net_io_counters", "net_connections", "net_if_addrs", # network + "net_if_stats", + "disk_io_counters", "disk_partitions", "disk_usage", # disk + # "sensors_temperatures", "sensors_battery", "sensors_fans" # sensors + "users", "boot_time", # others +] +# fmt: on + + +__all__.extend(_psplatform.__extra__all__) + +# Linux, FreeBSD +if hasattr(_psplatform.Process, "rlimit"): + # Populate global namespace with RLIM* constants. + _globals = globals() + _name = None + for _name in dir(_psplatform.cext): + if _name.startswith('RLIM') and _name.isupper(): + _globals[_name] = getattr(_psplatform.cext, _name) + __all__.append(_name) + del _globals, _name + +AF_LINK = _psplatform.AF_LINK + +__author__ = "Giampaolo Rodola'" +__version__ = "7.2.2" +version_info = tuple(int(num) for num in __version__.split('.')) + +_timer = getattr(time, 'monotonic', time.time) +_TOTAL_PHYMEM = None +_LOWEST_PID = None +_SENTINEL = object() + +# Sanity check in case the user messed up with psutil installation +# or did something weird with sys.path. In this case we might end +# up importing a python module using a C extension module which +# was compiled for a different version of psutil. +# We want to prevent that by failing sooner rather than later. +# See: https://github.com/giampaolo/psutil/issues/564 +if int(__version__.replace('.', '')) != getattr( + _psplatform.cext, 'version', None +): + msg = f"version conflict: {_psplatform.cext.__file__!r} C extension " + msg += "module was built for another version of psutil" + if hasattr(_psplatform.cext, 'version'): + v = ".".join(list(str(_psplatform.cext.version))) + msg += f" ({v} instead of {__version__})" + else: + msg += f" (different than {__version__})" + what = getattr( + _psplatform.cext, + "__file__", + "the existing psutil install directory", + ) + msg += f"; you may try to 'pip uninstall psutil', manually remove {what}" + msg += " or clean the virtual env somehow, then reinstall" + raise ImportError(msg) + + +# ===================================================================== +# --- Utils +# ===================================================================== + + +if hasattr(_psplatform, 'ppid_map'): + # Faster version (Windows and Linux). + _ppid_map = _psplatform.ppid_map +else: # pragma: no cover + + def _ppid_map(): + """Return a {pid: ppid, ...} dict for all running processes in + one shot. Used to speed up Process.children(). + """ + ret = {} + for pid in pids(): + try: + ret[pid] = _psplatform.Process(pid).ppid() + except (NoSuchProcess, ZombieProcess): + pass + return ret + + +def _pprint_secs(secs): + """Format seconds in a human readable form.""" + now = time.time() + secs_ago = int(now - secs) + fmt = "%H:%M:%S" if secs_ago < 60 * 60 * 24 else "%Y-%m-%d %H:%M:%S" + return datetime.datetime.fromtimestamp(secs).strftime(fmt) + + +def _check_conn_kind(kind): + """Check net_connections()'s `kind` parameter.""" + kinds = tuple(_common.conn_tmap) + if kind not in kinds: + msg = f"invalid kind argument {kind!r}; valid ones are: {kinds}" + raise ValueError(msg) + + +# ===================================================================== +# --- Process class +# ===================================================================== + + +class Process: + """Represents an OS process with the given PID. + If PID is omitted current process PID (os.getpid()) is used. + Raise NoSuchProcess if PID does not exist. + + Note that most of the methods of this class do not make sure that + the PID of the process being queried has been reused. That means + that you may end up retrieving information for another process. + + The only exceptions for which process identity is pre-emptively + checked and guaranteed are: + + - parent() + - children() + - nice() (set) + - ionice() (set) + - rlimit() (set) + - cpu_affinity (set) + - suspend() + - resume() + - send_signal() + - terminate() + - kill() + + To prevent this problem for all other methods you can use + is_running() before querying the process. + """ + + def __init__(self, pid=None): + self._init(pid) + + def _init(self, pid, _ignore_nsp=False): + if pid is None: + pid = os.getpid() + else: + if pid < 0: + msg = f"pid must be a positive integer (got {pid})" + raise ValueError(msg) + try: + _psplatform.cext.check_pid_range(pid) + except OverflowError as err: + msg = "process PID out of range" + raise NoSuchProcess(pid, msg=msg) from err + + self._pid = pid + self._name = None + self._exe = None + self._create_time = None + self._gone = False + self._pid_reused = False + self._hash = None + self._lock = threading.RLock() + # used for caching on Windows only (on POSIX ppid may change) + self._ppid = None + # platform-specific modules define an _psplatform.Process + # implementation class + self._proc = _psplatform.Process(pid) + self._last_sys_cpu_times = None + self._last_proc_cpu_times = None + self._exitcode = _SENTINEL + self._ident = (self.pid, None) + try: + self._ident = self._get_ident() + except AccessDenied: + # This should happen on Windows only, since we use the fast + # create time method. AFAIK, on all other platforms we are + # able to get create time for all PIDs. + pass + except ZombieProcess: + # Zombies can still be queried by this class (although + # not always) and pids() return them so just go on. + pass + except NoSuchProcess: + if not _ignore_nsp: + msg = "process PID not found" + raise NoSuchProcess(pid, msg=msg) from None + self._gone = True + + def _get_ident(self): + """Return a (pid, uid) tuple which is supposed to identify a + Process instance univocally over time. The PID alone is not + enough, as it can be assigned to a new process after this one + terminates, so we add process creation time to the mix. We need + this in order to prevent killing the wrong process later on. + This is also known as PID reuse or PID recycling problem. + + The reliability of this strategy mostly depends on + create_time() precision, which is 0.01 secs on Linux. The + assumption is that, after a process terminates, the kernel + won't reuse the same PID after such a short period of time + (0.01 secs). Technically this is inherently racy, but + practically it should be good enough. + + NOTE: unreliable on FreeBSD and OpenBSD as ctime is subject to + system clock updates. + """ + + if WINDOWS: + # Use create_time() fast method in order to speedup + # `process_iter()`. This means we'll get AccessDenied for + # most ADMIN processes, but that's fine since it means + # we'll also get AccessDenied on kill(). + # https://github.com/giampaolo/psutil/issues/2366#issuecomment-2381646555 + self._create_time = self._proc.create_time(fast_only=True) + return (self.pid, self._create_time) + elif LINUX or NETBSD or OSX: + # Use 'monotonic' process starttime since boot to form unique + # process identity, since it is stable over changes to system + # time. + return (self.pid, self._proc.create_time(monotonic=True)) + else: + return (self.pid, self.create_time()) + + def __str__(self): + info = collections.OrderedDict() + info["pid"] = self.pid + if self._name: + info['name'] = self._name + with self.oneshot(): + if self._pid_reused: + info["status"] = "terminated + PID reused" + else: + try: + info["name"] = self.name() + info["status"] = self.status() + except ZombieProcess: + info["status"] = "zombie" + except NoSuchProcess: + info["status"] = "terminated" + except AccessDenied: + pass + + if self._exitcode not in {_SENTINEL, None}: + info["exitcode"] = self._exitcode + if self._create_time is not None: + info['started'] = _pprint_secs(self._create_time) + + return "{}.{}({})".format( + self.__class__.__module__, + self.__class__.__name__, + ", ".join([f"{k}={v!r}" for k, v in info.items()]), + ) + + __repr__ = __str__ + + def __eq__(self, other): + # Test for equality with another Process object based + # on PID and creation time. + if not isinstance(other, Process): + return NotImplemented + if OPENBSD or NETBSD or SUNOS: # pragma: no cover + # Zombie processes on Open/NetBSD/illumos/Solaris have a + # creation time of 0.0. This covers the case when a process + # started normally (so it has a ctime), then it turned into a + # zombie. It's important to do this because is_running() + # depends on __eq__. + pid1, ident1 = self._ident + pid2, ident2 = other._ident + if pid1 == pid2: + if ident1 and not ident2: + try: + return self.status() == STATUS_ZOMBIE + except Error: + pass + return self._ident == other._ident + + def __ne__(self, other): + return not self == other + + def __hash__(self): + if self._hash is None: + self._hash = hash(self._ident) + return self._hash + + def _raise_if_pid_reused(self): + """Raises NoSuchProcess in case process PID has been reused.""" + if self._pid_reused or (not self.is_running() and self._pid_reused): + # We may directly raise NSP in here already if PID is just + # not running, but I prefer NSP to be raised naturally by + # the actual Process API call. This way unit tests will tell + # us if the API is broken (aka don't raise NSP when it + # should). We also remain consistent with all other "get" + # APIs which don't use _raise_if_pid_reused(). + msg = "process no longer exists and its PID has been reused" + raise NoSuchProcess(self.pid, self._name, msg=msg) + + @property + def pid(self): + """The process PID.""" + return self._pid + + # --- utility methods + + @contextlib.contextmanager + def oneshot(self): + """Utility context manager which considerably speeds up the + retrieval of multiple process information at the same time. + + Internally different process info (e.g. name, ppid, uids, + gids, ...) may be fetched by using the same routine, but + only one information is returned and the others are discarded. + When using this context manager the internal routine is + executed once (in the example below on name()) and the + other info are cached. + + The cache is cleared when exiting the context manager block. + The advice is to use this every time you retrieve more than + one information about the process. If you're lucky, you'll + get a hell of a speedup. + + >>> import psutil + >>> p = psutil.Process() + >>> with p.oneshot(): + ... p.name() # collect multiple info + ... p.cpu_times() # return cached value + ... p.cpu_percent() # return cached value + ... p.create_time() # return cached value + ... + >>> + """ + with self._lock: + if hasattr(self, "_cache"): + # NOOP: this covers the use case where the user enters the + # context twice: + # + # >>> with p.oneshot(): + # ... with p.oneshot(): + # ... + # + # Also, since as_dict() internally uses oneshot() + # I expect that the code below will be a pretty common + # "mistake" that the user will make, so let's guard + # against that: + # + # >>> with p.oneshot(): + # ... p.as_dict() + # ... + yield + else: + try: + # cached in case cpu_percent() is used + self.cpu_times.cache_activate(self) + # cached in case memory_percent() is used + self.memory_info.cache_activate(self) + # cached in case parent() is used + self.ppid.cache_activate(self) + # cached in case username() is used + if POSIX: + self.uids.cache_activate(self) + # specific implementation cache + self._proc.oneshot_enter() + yield + finally: + self.cpu_times.cache_deactivate(self) + self.memory_info.cache_deactivate(self) + self.ppid.cache_deactivate(self) + if POSIX: + self.uids.cache_deactivate(self) + self._proc.oneshot_exit() + + def as_dict(self, attrs=None, ad_value=None): + """Utility method returning process information as a + hashable dictionary. + If *attrs* is specified it must be a list of strings + reflecting available Process class' attribute names + (e.g. ['cpu_times', 'name']) else all public (read + only) attributes are assumed. + *ad_value* is the value which gets assigned in case + AccessDenied or ZombieProcess exception is raised when + retrieving that particular process information. + """ + valid_names = _as_dict_attrnames + if attrs is not None: + if not isinstance(attrs, (list, tuple, set, frozenset)): + msg = f"invalid attrs type {type(attrs)}" + raise TypeError(msg) + attrs = set(attrs) + invalid_names = attrs - valid_names + if invalid_names: + msg = "invalid attr name{} {}".format( + "s" if len(invalid_names) > 1 else "", + ", ".join(map(repr, invalid_names)), + ) + raise ValueError(msg) + + retdict = {} + ls = attrs or valid_names + with self.oneshot(): + for name in ls: + try: + if name == 'pid': + ret = self.pid + else: + meth = getattr(self, name) + ret = meth() + except (AccessDenied, ZombieProcess): + ret = ad_value + except NotImplementedError: + # in case of not implemented functionality (may happen + # on old or exotic systems) we want to crash only if + # the user explicitly asked for that particular attr + if attrs: + raise + continue + retdict[name] = ret + return retdict + + def parent(self): + """Return the parent process as a Process object pre-emptively + checking whether PID has been reused. + If no parent is known return None. + """ + lowest_pid = _LOWEST_PID if _LOWEST_PID is not None else pids()[0] + if self.pid == lowest_pid: + return None + ppid = self.ppid() + if ppid is not None: + # Get a fresh (non-cached) ctime in case the system clock + # was updated. TODO: use a monotonic ctime on platforms + # where it's supported. + proc_ctime = Process(self.pid).create_time() + try: + parent = Process(ppid) + if parent.create_time() <= proc_ctime: + return parent + # ...else ppid has been reused by another process + except NoSuchProcess: + pass + + def parents(self): + """Return the parents of this process as a list of Process + instances. If no parents are known return an empty list. + """ + parents = [] + proc = self.parent() + while proc is not None: + parents.append(proc) + proc = proc.parent() + return parents + + def is_running(self): + """Return whether this process is running. + + It also checks if PID has been reused by another process, in + which case it will remove the process from `process_iter()` + internal cache and return False. + """ + if self._gone or self._pid_reused: + return False + try: + # Checking if PID is alive is not enough as the PID might + # have been reused by another process. Process identity / + # uniqueness over time is guaranteed by (PID + creation + # time) and that is verified in __eq__. + self._pid_reused = self != Process(self.pid) + if self._pid_reused: + _pids_reused.add(self.pid) + raise NoSuchProcess(self.pid) + return True + except ZombieProcess: + # We should never get here as it's already handled in + # Process.__init__; here just for extra safety. + return True + except NoSuchProcess: + self._gone = True + return False + + # --- actual API + + @memoize_when_activated + def ppid(self): + """The process parent PID. + On Windows the return value is cached after first call. + """ + # On POSIX we don't want to cache the ppid as it may unexpectedly + # change to 1 (init) in case this process turns into a zombie: + # https://github.com/giampaolo/psutil/issues/321 + # http://stackoverflow.com/questions/356722/ + + # XXX should we check creation time here rather than in + # Process.parent()? + self._raise_if_pid_reused() + if POSIX: + return self._proc.ppid() + else: # pragma: no cover + self._ppid = self._ppid or self._proc.ppid() + return self._ppid + + def name(self): + """The process name. The return value is cached after first call.""" + # Process name is only cached on Windows as on POSIX it may + # change, see: + # https://github.com/giampaolo/psutil/issues/692 + if WINDOWS and self._name is not None: + return self._name + name = self._proc.name() + if POSIX and len(name) >= 15: + # On UNIX the name gets truncated to the first 15 characters. + # If it matches the first part of the cmdline we return that + # one instead because it's usually more explicative. + # Examples are "gnome-keyring-d" vs. "gnome-keyring-daemon". + try: + cmdline = self.cmdline() + except (AccessDenied, ZombieProcess): + # Just pass and return the truncated name: it's better + # than nothing. Note: there are actual cases where a + # zombie process can return a name() but not a + # cmdline(), see: + # https://github.com/giampaolo/psutil/issues/2239 + pass + else: + if cmdline: + extended_name = os.path.basename(cmdline[0]) + if extended_name.startswith(name): + name = extended_name + self._name = name + self._proc._name = name + return name + + def exe(self): + """The process executable as an absolute path. + May also be an empty string. + The return value is cached after first call. + """ + + def guess_it(fallback): + # try to guess exe from cmdline[0] in absence of a native + # exe representation + cmdline = self.cmdline() + if cmdline and hasattr(os, 'access') and hasattr(os, 'X_OK'): + exe = cmdline[0] # the possible exe + # Attempt to guess only in case of an absolute path. + # It is not safe otherwise as the process might have + # changed cwd. + if ( + os.path.isabs(exe) + and os.path.isfile(exe) + and os.access(exe, os.X_OK) + ): + return exe + if isinstance(fallback, AccessDenied): + raise fallback + return fallback + + if self._exe is None: + try: + exe = self._proc.exe() + except AccessDenied as err: + return guess_it(fallback=err) + else: + if not exe: + # underlying implementation can legitimately return an + # empty string; if that's the case we don't want to + # raise AD while guessing from the cmdline + try: + exe = guess_it(fallback=exe) + except AccessDenied: + pass + self._exe = exe + return self._exe + + def cmdline(self): + """The command line this process has been called with.""" + return self._proc.cmdline() + + def status(self): + """The process current status as a STATUS_* constant.""" + try: + return self._proc.status() + except ZombieProcess: + return STATUS_ZOMBIE + + def username(self): + """The name of the user that owns the process. + On UNIX this is calculated by using *real* process uid. + """ + if POSIX: + if pwd is None: + # might happen if python was installed from sources + msg = "requires pwd module shipped with standard python" + raise ImportError(msg) + real_uid = self.uids().real + try: + return pwd.getpwuid(real_uid).pw_name + except KeyError: + # the uid can't be resolved by the system + return str(real_uid) + else: + return self._proc.username() + + def create_time(self): + """The process creation time as a floating point number + expressed in seconds since the epoch (seconds since January 1, + 1970, at midnight UTC). The return value, which is cached after + first call, is based on the system clock, which means it may be + affected by changes such as manual adjustments or time + synchronization (e.g. NTP). + """ + if self._create_time is None: + self._create_time = self._proc.create_time() + return self._create_time + + def cwd(self): + """Process current working directory as an absolute path.""" + return self._proc.cwd() + + def nice(self, value=None): + """Get or set process niceness (priority).""" + if value is None: + return self._proc.nice_get() + else: + self._raise_if_pid_reused() + self._proc.nice_set(value) + + if POSIX: + + @memoize_when_activated + def uids(self): + """Return process UIDs as a (real, effective, saved) + namedtuple. + """ + return self._proc.uids() + + def gids(self): + """Return process GIDs as a (real, effective, saved) + namedtuple. + """ + return self._proc.gids() + + def terminal(self): + """The terminal associated with this process, if any, + else None. + """ + return self._proc.terminal() + + def num_fds(self): + """Return the number of file descriptors opened by this + process (POSIX only). + """ + return self._proc.num_fds() + + # Linux, BSD, AIX and Windows only + if hasattr(_psplatform.Process, "io_counters"): + + def io_counters(self): + """Return process I/O statistics as a + (read_count, write_count, read_bytes, write_bytes) + namedtuple. + Those are the number of read/write calls performed and the + amount of bytes read and written by the process. + """ + return self._proc.io_counters() + + # Linux and Windows + if hasattr(_psplatform.Process, "ionice_get"): + + def ionice(self, ioclass=None, value=None): + """Get or set process I/O niceness (priority). + + On Linux *ioclass* is one of the IOPRIO_CLASS_* constants. + *value* is a number which goes from 0 to 7. The higher the + value, the lower the I/O priority of the process. + + On Windows only *ioclass* is used and it can be set to 2 + (normal), 1 (low) or 0 (very low). + + Available on Linux and Windows > Vista only. + """ + if ioclass is None: + if value is not None: + msg = "'ioclass' argument must be specified" + raise ValueError(msg) + return self._proc.ionice_get() + else: + self._raise_if_pid_reused() + return self._proc.ionice_set(ioclass, value) + + # Linux / FreeBSD only + if hasattr(_psplatform.Process, "rlimit"): + + def rlimit(self, resource, limits=None): + """Get or set process resource limits as a (soft, hard) + tuple. + + *resource* is one of the RLIMIT_* constants. + *limits* is supposed to be a (soft, hard) tuple. + + See "man prlimit" for further info. + Available on Linux and FreeBSD only. + """ + if limits is not None: + self._raise_if_pid_reused() + return self._proc.rlimit(resource, limits) + + # Windows, Linux and FreeBSD only + if hasattr(_psplatform.Process, "cpu_affinity_get"): + + def cpu_affinity(self, cpus=None): + """Get or set process CPU affinity. + If specified, *cpus* must be a list of CPUs for which you + want to set the affinity (e.g. [0, 1]). + If an empty list is passed, all egible CPUs are assumed + (and set). + (Windows, Linux and BSD only). + """ + if cpus is None: + return sorted(set(self._proc.cpu_affinity_get())) + else: + self._raise_if_pid_reused() + if not cpus: + if hasattr(self._proc, "_get_eligible_cpus"): + cpus = self._proc._get_eligible_cpus() + else: + cpus = tuple(range(len(cpu_times(percpu=True)))) + self._proc.cpu_affinity_set(list(set(cpus))) + + # Linux, FreeBSD, SunOS + if hasattr(_psplatform.Process, "cpu_num"): + + def cpu_num(self): + """Return what CPU this process is currently running on. + The returned number should be <= psutil.cpu_count() + and <= len(psutil.cpu_percent(percpu=True)). + It may be used in conjunction with + psutil.cpu_percent(percpu=True) to observe the system + workload distributed across CPUs. + """ + return self._proc.cpu_num() + + # All platforms has it, but maybe not in the future. + if hasattr(_psplatform.Process, "environ"): + + def environ(self): + """The environment variables of the process as a dict. Note: this + might not reflect changes made after the process started. + """ + return self._proc.environ() + + if WINDOWS: + + def num_handles(self): + """Return the number of handles opened by this process + (Windows only). + """ + return self._proc.num_handles() + + def num_ctx_switches(self): + """Return the number of voluntary and involuntary context + switches performed by this process. + """ + return self._proc.num_ctx_switches() + + def num_threads(self): + """Return the number of threads used by this process.""" + return self._proc.num_threads() + + if hasattr(_psplatform.Process, "threads"): + + def threads(self): + """Return threads opened by process as a list of + (id, user_time, system_time) namedtuples representing + thread id and thread CPU times (user/system). + On OpenBSD this method requires root access. + """ + return self._proc.threads() + + def children(self, recursive=False): + """Return the children of this process as a list of Process + instances, pre-emptively checking whether PID has been reused. + If *recursive* is True return all the parent descendants. + + Example (A == this process): + + A ─┐ + │ + ├─ B (child) ─┐ + │ └─ X (grandchild) ─┐ + │ └─ Y (great grandchild) + ├─ C (child) + └─ D (child) + + >>> import psutil + >>> p = psutil.Process() + >>> p.children() + B, C, D + >>> p.children(recursive=True) + B, X, Y, C, D + + Note that in the example above if process X disappears + process Y won't be listed as the reference to process A + is lost. + """ + self._raise_if_pid_reused() + ppid_map = _ppid_map() + # Get a fresh (non-cached) ctime in case the system clock was + # updated. TODO: use a monotonic ctime on platforms where it's + # supported. + proc_ctime = Process(self.pid).create_time() + ret = [] + if not recursive: + for pid, ppid in ppid_map.items(): + if ppid == self.pid: + try: + child = Process(pid) + # if child happens to be older than its parent + # (self) it means child's PID has been reused + if proc_ctime <= child.create_time(): + ret.append(child) + except (NoSuchProcess, ZombieProcess): + pass + else: + # Construct a {pid: [child pids]} dict + reverse_ppid_map = collections.defaultdict(list) + for pid, ppid in ppid_map.items(): + reverse_ppid_map[ppid].append(pid) + # Recursively traverse that dict, starting from self.pid, + # such that we only call Process() on actual children + seen = set() + stack = [self.pid] + while stack: + pid = stack.pop() + if pid in seen: + # Since pids can be reused while the ppid_map is + # constructed, there may be rare instances where + # there's a cycle in the recorded process "tree". + continue + seen.add(pid) + for child_pid in reverse_ppid_map[pid]: + try: + child = Process(child_pid) + # if child happens to be older than its parent + # (self) it means child's PID has been reused + intime = proc_ctime <= child.create_time() + if intime: + ret.append(child) + stack.append(child_pid) + except (NoSuchProcess, ZombieProcess): + pass + return ret + + def cpu_percent(self, interval=None): + """Return a float representing the current process CPU + utilization as a percentage. + + When *interval* is 0.0 or None (default) compares process times + to system CPU times elapsed since last call, returning + immediately (non-blocking). That means that the first time + this is called it will return a meaningful 0.0 value. + + When *interval* is > 0.0 compares process times to system CPU + times elapsed before and after the interval (blocking). + + In this case is recommended for accuracy that this function + be called with at least 0.1 seconds between calls. + + A value > 100.0 can be returned in case of processes running + multiple threads on different CPU cores. + + The returned value is explicitly NOT split evenly between + all available logical CPUs. This means that a busy loop process + running on a system with 2 logical CPUs will be reported as + having 100% CPU utilization instead of 50%. + + Examples: + + >>> import psutil + >>> p = psutil.Process(os.getpid()) + >>> # blocking + >>> p.cpu_percent(interval=1) + 2.0 + >>> # non-blocking (percentage since last call) + >>> p.cpu_percent(interval=None) + 2.9 + >>> + """ + blocking = interval is not None and interval > 0.0 + if interval is not None and interval < 0: + msg = f"interval is not positive (got {interval!r})" + raise ValueError(msg) + num_cpus = cpu_count() or 1 + + def timer(): + return _timer() * num_cpus + + if blocking: + st1 = timer() + pt1 = self._proc.cpu_times() + time.sleep(interval) + st2 = timer() + pt2 = self._proc.cpu_times() + else: + st1 = self._last_sys_cpu_times + pt1 = self._last_proc_cpu_times + st2 = timer() + pt2 = self._proc.cpu_times() + if st1 is None or pt1 is None: + self._last_sys_cpu_times = st2 + self._last_proc_cpu_times = pt2 + return 0.0 + + delta_proc = (pt2.user - pt1.user) + (pt2.system - pt1.system) + delta_time = st2 - st1 + # reset values for next call in case of interval == None + self._last_sys_cpu_times = st2 + self._last_proc_cpu_times = pt2 + + try: + # This is the utilization split evenly between all CPUs. + # E.g. a busy loop process on a 2-CPU-cores system at this + # point is reported as 50% instead of 100%. + overall_cpus_percent = (delta_proc / delta_time) * 100 + except ZeroDivisionError: + # interval was too low + return 0.0 + else: + # Note 1: + # in order to emulate "top" we multiply the value for the num + # of CPU cores. This way the busy process will be reported as + # having 100% (or more) usage. + # + # Note 2: + # taskmgr.exe on Windows differs in that it will show 50% + # instead. + # + # Note 3: + # a percentage > 100 is legitimate as it can result from a + # process with multiple threads running on different CPU + # cores (top does the same), see: + # http://stackoverflow.com/questions/1032357 + # https://github.com/giampaolo/psutil/issues/474 + single_cpu_percent = overall_cpus_percent * num_cpus + return round(single_cpu_percent, 1) + + @memoize_when_activated + def cpu_times(self): + """Return a (user, system, children_user, children_system) + namedtuple representing the accumulated process time, in + seconds. + This is similar to os.times() but per-process. + On macOS and Windows children_user and children_system are + always set to 0. + """ + return self._proc.cpu_times() + + @memoize_when_activated + def memory_info(self): + """Return a namedtuple with variable fields depending on the + platform, representing memory information about the process. + + The "portable" fields available on all platforms are `rss` and `vms`. + + All numbers are expressed in bytes. + """ + return self._proc.memory_info() + + def memory_full_info(self): + """This method returns the same information as memory_info(), + plus, on some platform (Linux, macOS, Windows), also provides + additional metrics (USS, PSS and swap). + The additional metrics provide a better representation of actual + process memory usage. + + Namely USS is the memory which is unique to a process and which + would be freed if the process was terminated right now. + + It does so by passing through the whole process address. + As such it usually requires higher user privileges than + memory_info() and is considerably slower. + """ + return self._proc.memory_full_info() + + def memory_percent(self, memtype="rss"): + """Compare process memory to total physical system memory and + calculate process memory utilization as a percentage. + *memtype* argument is a string that dictates what type of + process memory you want to compare against (defaults to "rss"). + The list of available strings can be obtained like this: + + >>> psutil.Process().memory_info()._fields + ('rss', 'vms', 'shared', 'text', 'lib', 'data', 'dirty', 'uss', 'pss') + """ + valid_types = list(_ntp.pfullmem._fields) + if memtype not in valid_types: + msg = ( + f"invalid memtype {memtype!r}; valid types are" + f" {tuple(valid_types)!r}" + ) + raise ValueError(msg) + fun = ( + self.memory_info + if memtype in _ntp.pmem._fields + else self.memory_full_info + ) + metrics = fun() + value = getattr(metrics, memtype) + + # use cached value if available + total_phymem = _TOTAL_PHYMEM or virtual_memory().total + if not total_phymem > 0: + # we should never get here + msg = ( + "can't calculate process memory percent because total physical" + f" system memory is not positive ({total_phymem!r})" + ) + raise ValueError(msg) + return (value / float(total_phymem)) * 100 + + if hasattr(_psplatform.Process, "memory_maps"): + + def memory_maps(self, grouped=True): + """Return process' mapped memory regions as a list of namedtuples + whose fields are variable depending on the platform. + + If *grouped* is True the mapped regions with the same 'path' + are grouped together and the different memory fields are summed. + + If *grouped* is False every mapped region is shown as a single + entity and the namedtuple will also include the mapped region's + address space ('addr') and permission set ('perms'). + """ + it = self._proc.memory_maps() + if grouped: + d = {} + for tupl in it: + path = tupl[2] + nums = tupl[3:] + try: + d[path] = list(map(lambda x, y: x + y, d[path], nums)) + except KeyError: + d[path] = nums + return [_ntp.pmmap_grouped(path, *d[path]) for path in d] + else: + return [_ntp.pmmap_ext(*x) for x in it] + + def open_files(self): + """Return files opened by process as a list of + (path, fd) namedtuples including the absolute file name + and file descriptor number. + """ + return self._proc.open_files() + + def net_connections(self, kind='inet'): + """Return socket connections opened by process as a list of + (fd, family, type, laddr, raddr, status) namedtuples. + The *kind* parameter filters for connections that match the + following criteria: + + +------------+----------------------------------------------------+ + | Kind Value | Connections using | + +------------+----------------------------------------------------+ + | inet | IPv4 and IPv6 | + | inet4 | IPv4 | + | inet6 | IPv6 | + | tcp | TCP | + | tcp4 | TCP over IPv4 | + | tcp6 | TCP over IPv6 | + | udp | UDP | + | udp4 | UDP over IPv4 | + | udp6 | UDP over IPv6 | + | unix | UNIX socket (both UDP and TCP protocols) | + | all | the sum of all the possible families and protocols | + +------------+----------------------------------------------------+ + """ + _check_conn_kind(kind) + return self._proc.net_connections(kind) + + @_common.deprecated_method(replacement="net_connections") + def connections(self, kind="inet"): + return self.net_connections(kind=kind) + + # --- signals + + if POSIX: + + def _send_signal(self, sig): + assert not self.pid < 0, self.pid + self._raise_if_pid_reused() + + pid, ppid, name = self.pid, self._ppid, self._name + if pid == 0: + # see "man 2 kill" + msg = ( + "preventing sending signal to process with PID 0 as it " + "would affect every process in the process group of the " + "calling process (os.getpid()) instead of PID 0" + ) + raise ValueError(msg) + try: + os.kill(pid, sig) + except ProcessLookupError as err: + if OPENBSD and pid_exists(pid): + # We do this because os.kill() lies in case of + # zombie processes. + raise ZombieProcess(pid, name, ppid) from err + self._gone = True + raise NoSuchProcess(pid, name) from err + except PermissionError as err: + raise AccessDenied(pid, name) from err + + def send_signal(self, sig): + """Send a signal *sig* to process pre-emptively checking + whether PID has been reused (see signal module constants) . + On Windows only SIGTERM is valid and is treated as an alias + for kill(). + """ + if POSIX: + self._send_signal(sig) + else: # pragma: no cover + self._raise_if_pid_reused() + if sig != signal.SIGTERM and not self.is_running(): + msg = "process no longer exists" + raise NoSuchProcess(self.pid, self._name, msg=msg) + self._proc.send_signal(sig) + + def suspend(self): + """Suspend process execution with SIGSTOP pre-emptively checking + whether PID has been reused. + On Windows this has the effect of suspending all process threads. + """ + if POSIX: + self._send_signal(signal.SIGSTOP) + else: # pragma: no cover + self._raise_if_pid_reused() + self._proc.suspend() + + def resume(self): + """Resume process execution with SIGCONT pre-emptively checking + whether PID has been reused. + On Windows this has the effect of resuming all process threads. + """ + if POSIX: + self._send_signal(signal.SIGCONT) + else: # pragma: no cover + self._raise_if_pid_reused() + self._proc.resume() + + def terminate(self): + """Terminate the process with SIGTERM pre-emptively checking + whether PID has been reused. + On Windows this is an alias for kill(). + """ + if POSIX: + self._send_signal(signal.SIGTERM) + else: # pragma: no cover + self._raise_if_pid_reused() + self._proc.kill() + + def kill(self): + """Kill the current process with SIGKILL pre-emptively checking + whether PID has been reused. + """ + if POSIX: + self._send_signal(signal.SIGKILL) + else: # pragma: no cover + self._raise_if_pid_reused() + self._proc.kill() + + def wait(self, timeout=None): + """Wait for process to terminate, and if process is a children + of os.getpid(), also return its exit code, else None. + On Windows there's no such limitation (exit code is always + returned). + + If the process is already terminated, immediately return None + instead of raising NoSuchProcess. + + If *timeout* (in seconds) is specified and process is still + alive, raise TimeoutExpired. + + If *timeout=0* either return immediately or raise + TimeoutExpired (non-blocking). + + To wait for multiple Process objects use psutil.wait_procs(). + """ + if self.pid == 0: + msg = "can't wait for PID 0" + raise ValueError(msg) + if timeout is not None: + if not isinstance(timeout, (int, float)): + msg = f"timeout must be an int or float (got {type(timeout)})" + raise TypeError(msg) + if timeout < 0: + msg = f"timeout must be positive or zero (got {timeout})" + raise ValueError(msg) + + if self._exitcode is not _SENTINEL: + return self._exitcode + + try: + self._exitcode = self._proc.wait(timeout) + except TimeoutExpired as err: + exc = TimeoutExpired(timeout, pid=self.pid, name=self._name) + raise exc from err + + return self._exitcode + + +# The valid attr names which can be processed by Process.as_dict(). +# fmt: off +_as_dict_attrnames = { + x for x in dir(Process) if not x.startswith("_") and x not in + {'send_signal', 'suspend', 'resume', 'terminate', 'kill', 'wait', + 'is_running', 'as_dict', 'parent', 'parents', 'children', 'rlimit', + 'connections', 'oneshot'} +} +# fmt: on + + +# ===================================================================== +# --- Popen class +# ===================================================================== + + +class Popen(Process): + """Same as subprocess.Popen, but in addition it provides all + psutil.Process methods in a single class. + For the following methods which are common to both classes, psutil + implementation takes precedence: + + * send_signal() + * terminate() + * kill() + + This is done in order to avoid killing another process in case its + PID has been reused, fixing BPO-6973. + + >>> import psutil + >>> from subprocess import PIPE + >>> p = psutil.Popen(["python", "-c", "print 'hi'"], stdout=PIPE) + >>> p.name() + 'python' + >>> p.uids() + user(real=1000, effective=1000, saved=1000) + >>> p.username() + 'giampaolo' + >>> p.communicate() + ('hi', None) + >>> p.terminate() + >>> p.wait(timeout=2) + 0 + >>> + """ + + def __init__(self, *args, **kwargs): + # Explicitly avoid to raise NoSuchProcess in case the process + # spawned by subprocess.Popen terminates too quickly, see: + # https://github.com/giampaolo/psutil/issues/193 + self.__subproc = subprocess.Popen(*args, **kwargs) + self._init(self.__subproc.pid, _ignore_nsp=True) + + def __dir__(self): + return sorted(set(dir(Popen) + dir(subprocess.Popen))) + + def __enter__(self): + if hasattr(self.__subproc, '__enter__'): + self.__subproc.__enter__() + return self + + def __exit__(self, *args, **kwargs): + if hasattr(self.__subproc, '__exit__'): + return self.__subproc.__exit__(*args, **kwargs) + else: + if self.stdout: + self.stdout.close() + if self.stderr: + self.stderr.close() + try: + # Flushing a BufferedWriter may raise an error. + if self.stdin: + self.stdin.close() + finally: + # Wait for the process to terminate, to avoid zombies. + self.wait() + + def __getattribute__(self, name): + try: + return object.__getattribute__(self, name) + except AttributeError: + try: + return object.__getattribute__(self.__subproc, name) + except AttributeError: + msg = f"{self.__class__!r} has no attribute {name!r}" + raise AttributeError(msg) from None + + def wait(self, timeout=None): + if self.__subproc.returncode is not None: + return self.__subproc.returncode + ret = super().wait(timeout) + self.__subproc.returncode = ret + return ret + + +# ===================================================================== +# --- system processes related functions +# ===================================================================== + + +def pids(): + """Return a list of current running PIDs.""" + global _LOWEST_PID + ret = sorted(_psplatform.pids()) + _LOWEST_PID = ret[0] + return ret + + +def pid_exists(pid): + """Return True if given PID exists in the current process list. + This is faster than doing "pid in psutil.pids()" and + should be preferred. + """ + if pid < 0: + return False + elif pid == 0 and POSIX: + # On POSIX we use os.kill() to determine PID existence. + # According to "man 2 kill" PID 0 has a special meaning + # though: it refers to <> and that is not we want + # to do here. + return pid in pids() + else: + return _psplatform.pid_exists(pid) + + +_pmap = {} +_pids_reused = set() + + +def process_iter(attrs=None, ad_value=None): + """Return a generator yielding a Process instance for all + running processes. + + Every new Process instance is only created once and then cached + into an internal table which is updated every time this is used. + Cache can optionally be cleared via `process_iter.cache_clear()`. + + The sorting order in which processes are yielded is based on + their PIDs. + + *attrs* and *ad_value* have the same meaning as in + Process.as_dict(). If *attrs* is specified as_dict() is called + and the resulting dict is stored as a 'info' attribute attached + to returned Process instance. + If *attrs* is an empty list it will retrieve all process info + (slow). + """ + global _pmap + + def add(pid): + proc = Process(pid) + pmap[proc.pid] = proc + return proc + + def remove(pid): + pmap.pop(pid, None) + + pmap = _pmap.copy() + a = set(pids()) + b = set(pmap.keys()) + new_pids = a - b + gone_pids = b - a + for pid in gone_pids: + remove(pid) + while _pids_reused: + pid = _pids_reused.pop() + debug(f"refreshing Process instance for reused PID {pid}") + remove(pid) + try: + ls = sorted(list(pmap.items()) + list(dict.fromkeys(new_pids).items())) + for pid, proc in ls: + try: + if proc is None: # new process + proc = add(pid) + if attrs is not None: + proc.info = proc.as_dict(attrs=attrs, ad_value=ad_value) + yield proc + except NoSuchProcess: + remove(pid) + finally: + _pmap = pmap + + +process_iter.cache_clear = lambda: _pmap.clear() # noqa: PLW0108 +process_iter.cache_clear.__doc__ = "Clear process_iter() internal cache." + + +def wait_procs(procs, timeout=None, callback=None): + """Convenience function which waits for a list of processes to + terminate. + + Return a (gone, alive) tuple indicating which processes + are gone and which ones are still alive. + + The gone ones will have a new *returncode* attribute indicating + process exit status (may be None). + + *callback* is a function which gets called every time a process + terminates (a Process instance is passed as callback argument). + + Function will return as soon as all processes terminate or when + *timeout* occurs. + Differently from Process.wait() it will not raise TimeoutExpired if + *timeout* occurs. + + Typical use case is: + + - send SIGTERM to a list of processes + - give them some time to terminate + - send SIGKILL to those ones which are still alive + + Example: + + >>> def on_terminate(proc): + ... print("process {} terminated".format(proc)) + ... + >>> for p in procs: + ... p.terminate() + ... + >>> gone, alive = wait_procs(procs, timeout=3, callback=on_terminate) + >>> for p in alive: + ... p.kill() + """ + + def check_gone(proc, timeout): + try: + returncode = proc.wait(timeout=timeout) + except (TimeoutExpired, subprocess.TimeoutExpired): + pass + else: + if returncode is not None or not proc.is_running(): + # Set new Process instance attribute. + proc.returncode = returncode + gone.add(proc) + if callback is not None: + callback(proc) + + if timeout is not None and not timeout >= 0: + msg = f"timeout must be a positive integer, got {timeout}" + raise ValueError(msg) + if callback is not None and not callable(callback): + msg = f"callback {callback!r} is not a callable" + raise TypeError(msg) + + gone = set() + alive = set(procs) + if timeout is not None: + deadline = _timer() + timeout + + while alive: + if timeout is not None and timeout <= 0: + break + for proc in alive: + # Make sure that every complete iteration (all processes) + # will last max 1 sec. + # We do this because we don't want to wait too long on a + # single process: in case it terminates too late other + # processes may disappear in the meantime and their PID + # reused. + max_timeout = 1.0 / len(alive) + if timeout is not None: + timeout = min((deadline - _timer()), max_timeout) + if timeout <= 0: + break + check_gone(proc, timeout) + else: + check_gone(proc, max_timeout) + alive = alive - gone # noqa: PLR6104 + + if alive: + # Last attempt over processes survived so far. + # timeout == 0 won't make this function wait any further. + for proc in alive: + check_gone(proc, 0) + alive = alive - gone # noqa: PLR6104 + + return (list(gone), list(alive)) + + +# ===================================================================== +# --- CPU related functions +# ===================================================================== + + +def cpu_count(logical=True): + """Return the number of logical CPUs in the system (same as + os.cpu_count()). + + If *logical* is False return the number of physical cores only + (e.g. hyper thread CPUs are excluded). + + Return None if undetermined. + + The return value is cached after first call. + If desired cache can be cleared like this: + + >>> psutil.cpu_count.cache_clear() + """ + if logical: + ret = _psplatform.cpu_count_logical() + else: + ret = _psplatform.cpu_count_cores() + if ret is not None and ret < 1: + ret = None + return ret + + +def cpu_times(percpu=False): + """Return system-wide CPU times as a namedtuple. + Every CPU time represents the seconds the CPU has spent in the + given mode. The namedtuple's fields availability varies depending on the + platform: + + - user + - system + - idle + - nice (UNIX) + - iowait (Linux) + - irq (Linux, FreeBSD) + - softirq (Linux) + - steal (Linux >= 2.6.11) + - guest (Linux >= 2.6.24) + - guest_nice (Linux >= 3.2.0) + + When *percpu* is True return a list of namedtuples for each CPU. + First element of the list refers to first CPU, second element + to second CPU and so on. + The order of the list is consistent across calls. + """ + if not percpu: + return _psplatform.cpu_times() + else: + return _psplatform.per_cpu_times() + + +try: + _last_cpu_times = {threading.current_thread().ident: cpu_times()} +except Exception: # noqa: BLE001 + # Don't want to crash at import time. + _last_cpu_times = {} + +try: + _last_per_cpu_times = { + threading.current_thread().ident: cpu_times(percpu=True) + } +except Exception: # noqa: BLE001 + # Don't want to crash at import time. + _last_per_cpu_times = {} + + +def _cpu_tot_time(times): + """Given a cpu_time() ntuple calculates the total CPU time + (including idle time). + """ + tot = sum(times) + if LINUX: + # On Linux guest times are already accounted in "user" or + # "nice" times, so we subtract them from total. + # Htop does the same. References: + # https://github.com/giampaolo/psutil/pull/940 + # http://unix.stackexchange.com/questions/178045 + # https://github.com/torvalds/linux/blob/ + # 447976ef4fd09b1be88b316d1a81553f1aa7cd07/kernel/sched/ + # cputime.c#L158 + tot -= getattr(times, "guest", 0) # Linux 2.6.24+ + tot -= getattr(times, "guest_nice", 0) # Linux 3.2.0+ + return tot + + +def _cpu_busy_time(times): + """Given a cpu_time() ntuple calculates the busy CPU time. + We do so by subtracting all idle CPU times. + """ + busy = _cpu_tot_time(times) + busy -= times.idle + # Linux: "iowait" is time during which the CPU does not do anything + # (waits for IO to complete). On Linux IO wait is *not* accounted + # in "idle" time so we subtract it. Htop does the same. + # References: + # https://github.com/torvalds/linux/blob/ + # 447976ef4fd09b1be88b316d1a81553f1aa7cd07/kernel/sched/cputime.c#L244 + busy -= getattr(times, "iowait", 0) + return busy + + +def _cpu_times_deltas(t1, t2): + assert t1._fields == t2._fields, (t1, t2) + field_deltas = [] + for field in _ntp.scputimes._fields: + field_delta = getattr(t2, field) - getattr(t1, field) + # CPU times are always supposed to increase over time + # or at least remain the same and that's because time + # cannot go backwards. + # Surprisingly sometimes this might not be the case (at + # least on Windows and Linux), see: + # https://github.com/giampaolo/psutil/issues/392 + # https://github.com/giampaolo/psutil/issues/645 + # https://github.com/giampaolo/psutil/issues/1210 + # Trim negative deltas to zero to ignore decreasing fields. + # top does the same. Reference: + # https://gitlab.com/procps-ng/procps/blob/v3.3.12/top/top.c#L5063 + field_delta = max(0, field_delta) + field_deltas.append(field_delta) + return _ntp.scputimes(*field_deltas) + + +def cpu_percent(interval=None, percpu=False): + """Return a float representing the current system-wide CPU + utilization as a percentage. + + When *interval* is > 0.0 compares system CPU times elapsed before + and after the interval (blocking). + + When *interval* is 0.0 or None compares system CPU times elapsed + since last call or module import, returning immediately (non + blocking). That means the first time this is called it will + return a meaningless 0.0 value which you should ignore. + In this case is recommended for accuracy that this function be + called with at least 0.1 seconds between calls. + + When *percpu* is True returns a list of floats representing the + utilization as a percentage for each CPU. + First element of the list refers to first CPU, second element + to second CPU and so on. + The order of the list is consistent across calls. + + Examples: + + >>> # blocking, system-wide + >>> psutil.cpu_percent(interval=1) + 2.0 + >>> + >>> # blocking, per-cpu + >>> psutil.cpu_percent(interval=1, percpu=True) + [2.0, 1.0] + >>> + >>> # non-blocking (percentage since last call) + >>> psutil.cpu_percent(interval=None) + 2.9 + >>> + """ + tid = threading.current_thread().ident + blocking = interval is not None and interval > 0.0 + if interval is not None and interval < 0: + msg = f"interval is not positive (got {interval})" + raise ValueError(msg) + + def calculate(t1, t2): + times_delta = _cpu_times_deltas(t1, t2) + all_delta = _cpu_tot_time(times_delta) + busy_delta = _cpu_busy_time(times_delta) + + try: + busy_perc = (busy_delta / all_delta) * 100 + except ZeroDivisionError: + return 0.0 + else: + return round(busy_perc, 1) + + # system-wide usage + if not percpu: + if blocking: + t1 = cpu_times() + time.sleep(interval) + else: + t1 = _last_cpu_times.get(tid) or cpu_times() + _last_cpu_times[tid] = cpu_times() + return calculate(t1, _last_cpu_times[tid]) + # per-cpu usage + else: + ret = [] + if blocking: + tot1 = cpu_times(percpu=True) + time.sleep(interval) + else: + tot1 = _last_per_cpu_times.get(tid) or cpu_times(percpu=True) + _last_per_cpu_times[tid] = cpu_times(percpu=True) + for t1, t2 in zip(tot1, _last_per_cpu_times[tid]): + ret.append(calculate(t1, t2)) + return ret + + +# Use a separate dict for cpu_times_percent(), so it's independent from +# cpu_percent() and they can both be used within the same program. +_last_cpu_times_2 = _last_cpu_times.copy() +_last_per_cpu_times_2 = _last_per_cpu_times.copy() + + +def cpu_times_percent(interval=None, percpu=False): + """Same as cpu_percent() but provides utilization percentages + for each specific CPU time as is returned by cpu_times(). + For instance, on Linux we'll get: + + >>> cpu_times_percent() + cpupercent(user=4.8, nice=0.0, system=4.8, idle=90.5, iowait=0.0, + irq=0.0, softirq=0.0, steal=0.0, guest=0.0, guest_nice=0.0) + >>> + + *interval* and *percpu* arguments have the same meaning as in + cpu_percent(). + """ + tid = threading.current_thread().ident + blocking = interval is not None and interval > 0.0 + if interval is not None and interval < 0: + msg = f"interval is not positive (got {interval!r})" + raise ValueError(msg) + + def calculate(t1, t2): + nums = [] + times_delta = _cpu_times_deltas(t1, t2) + all_delta = _cpu_tot_time(times_delta) + # "scale" is the value to multiply each delta with to get percentages. + # We use "max" to avoid division by zero (if all_delta is 0, then all + # fields are 0 so percentages will be 0 too. all_delta cannot be a + # fraction because cpu times are integers) + scale = 100.0 / max(1, all_delta) + for field_delta in times_delta: + field_perc = field_delta * scale + field_perc = round(field_perc, 1) + # make sure we don't return negative values or values over 100% + field_perc = min(max(0.0, field_perc), 100.0) + nums.append(field_perc) + return _ntp.scputimes(*nums) + + # system-wide usage + if not percpu: + if blocking: + t1 = cpu_times() + time.sleep(interval) + else: + t1 = _last_cpu_times_2.get(tid) or cpu_times() + _last_cpu_times_2[tid] = cpu_times() + return calculate(t1, _last_cpu_times_2[tid]) + # per-cpu usage + else: + ret = [] + if blocking: + tot1 = cpu_times(percpu=True) + time.sleep(interval) + else: + tot1 = _last_per_cpu_times_2.get(tid) or cpu_times(percpu=True) + _last_per_cpu_times_2[tid] = cpu_times(percpu=True) + for t1, t2 in zip(tot1, _last_per_cpu_times_2[tid]): + ret.append(calculate(t1, t2)) + return ret + + +def cpu_stats(): + """Return CPU statistics.""" + return _psplatform.cpu_stats() + + +if hasattr(_psplatform, "cpu_freq"): + + def cpu_freq(percpu=False): + """Return CPU frequency as a namedtuple including current, + min and max frequency expressed in Mhz. + + If *percpu* is True and the system supports per-cpu frequency + retrieval (Linux only) a list of frequencies is returned for + each CPU. If not a list with one element is returned. + """ + ret = _psplatform.cpu_freq() + if percpu: + return ret + else: + num_cpus = float(len(ret)) + if num_cpus == 0: + return None + elif num_cpus == 1: + return ret[0] + else: + currs, mins, maxs = 0.0, 0.0, 0.0 + set_none = False + for cpu in ret: + currs += cpu.current + # On Linux if /proc/cpuinfo is used min/max are set + # to None. + if LINUX and cpu.min is None: + set_none = True + continue + mins += cpu.min + maxs += cpu.max + + current = currs / num_cpus + + if set_none: + min_ = max_ = None + else: + min_ = mins / num_cpus + max_ = maxs / num_cpus + + return _ntp.scpufreq(current, min_, max_) + + __all__.append("cpu_freq") + + +if hasattr(os, "getloadavg") or hasattr(_psplatform, "getloadavg"): + # Perform this hasattr check once on import time to either use the + # platform based code or proxy straight from the os module. + if hasattr(os, "getloadavg"): + getloadavg = os.getloadavg + else: + getloadavg = _psplatform.getloadavg + + __all__.append("getloadavg") + + +# ===================================================================== +# --- system memory related functions +# ===================================================================== + + +def virtual_memory(): + """Return statistics about system memory usage as a namedtuple + including the following fields, expressed in bytes: + + - total: + total physical memory available. + + - available: + the memory that can be given instantly to processes without the + system going into swap. + This is calculated by summing different memory values depending + on the platform and it is supposed to be used to monitor actual + memory usage in a cross platform fashion. + + - percent: + the percentage usage calculated as (total - available) / total * 100 + + - used: + memory used, calculated differently depending on the platform and + designed for informational purposes only: + macOS: active + wired + BSD: active + wired + cached + Linux: total - free + + - free: + memory not being used at all (zeroed) that is readily available; + note that this doesn't reflect the actual memory available + (use 'available' instead) + + Platform-specific fields: + + - active (UNIX): + memory currently in use or very recently used, and so it is in RAM. + + - inactive (UNIX): + memory that is marked as not used. + + - buffers (BSD, Linux): + cache for things like file system metadata. + + - cached (BSD, macOS): + cache for various things. + + - wired (macOS, BSD): + memory that is marked to always stay in RAM. It is never moved to disk. + + - shared (BSD): + memory that may be simultaneously accessed by multiple processes. + + The sum of 'used' and 'available' does not necessarily equal total. + On Windows 'available' and 'free' are the same. + """ + global _TOTAL_PHYMEM + ret = _psplatform.virtual_memory() + # cached for later use in Process.memory_percent() + _TOTAL_PHYMEM = ret.total + return ret + + +def swap_memory(): + """Return system swap memory statistics as a namedtuple including + the following fields: + + - total: total swap memory in bytes + - used: used swap memory in bytes + - free: free swap memory in bytes + - percent: the percentage usage + - sin: no. of bytes the system has swapped in from disk (cumulative) + - sout: no. of bytes the system has swapped out from disk (cumulative) + + 'sin' and 'sout' on Windows are meaningless and always set to 0. + """ + return _psplatform.swap_memory() + + +# ===================================================================== +# --- disks/partitions related functions +# ===================================================================== + + +def disk_usage(path): + """Return disk usage statistics about the given *path* as a + namedtuple including total, used and free space expressed in bytes + plus the percentage usage. + """ + return _psplatform.disk_usage(path) + + +def disk_partitions(all=False): + """Return mounted partitions as a list of + (device, mountpoint, fstype, opts) namedtuple. + 'opts' field is a raw string separated by commas indicating mount + options which may vary depending on the platform. + + If *all* parameter is False return physical devices only and ignore + all others. + """ + return _psplatform.disk_partitions(all) + + +def disk_io_counters(perdisk=False, nowrap=True): + """Return system disk I/O statistics as a namedtuple including + the following fields: + + - read_count: number of reads + - write_count: number of writes + - read_bytes: number of bytes read + - write_bytes: number of bytes written + - read_time: time spent reading from disk (in ms) + - write_time: time spent writing to disk (in ms) + + Platform specific: + + - busy_time: (Linux, FreeBSD) time spent doing actual I/Os (in ms) + - read_merged_count (Linux): number of merged reads + - write_merged_count (Linux): number of merged writes + + If *perdisk* is True return the same information for every + physical disk installed on the system as a dictionary + with partition names as the keys and the namedtuple + described above as the values. + + If *nowrap* is True it detects and adjust the numbers which overflow + and wrap (restart from 0) and add "old value" to "new value" so that + the returned numbers will always be increasing or remain the same, + but never decrease. + "disk_io_counters.cache_clear()" can be used to invalidate the + cache. + + On recent Windows versions 'diskperf -y' command may need to be + executed first otherwise this function won't find any disk. + """ + kwargs = dict(perdisk=perdisk) if LINUX else {} + rawdict = _psplatform.disk_io_counters(**kwargs) + if not rawdict: + return {} if perdisk else None + if nowrap: + rawdict = _wrap_numbers(rawdict, 'psutil.disk_io_counters') + if perdisk: + for disk, fields in rawdict.items(): + rawdict[disk] = _ntp.sdiskio(*fields) + return rawdict + else: + return _ntp.sdiskio(*(sum(x) for x in zip(*rawdict.values()))) + + +disk_io_counters.cache_clear = functools.partial( + _wrap_numbers.cache_clear, 'psutil.disk_io_counters' +) +disk_io_counters.cache_clear.__doc__ = "Clears nowrap argument cache" + + +# ===================================================================== +# --- network related functions +# ===================================================================== + + +def net_io_counters(pernic=False, nowrap=True): + """Return network I/O statistics as a namedtuple including + the following fields: + + - bytes_sent: number of bytes sent + - bytes_recv: number of bytes received + - packets_sent: number of packets sent + - packets_recv: number of packets received + - errin: total number of errors while receiving + - errout: total number of errors while sending + - dropin: total number of incoming packets which were dropped + - dropout: total number of outgoing packets which were dropped + (always 0 on macOS and BSD) + + If *pernic* is True return the same information for every + network interface installed on the system as a dictionary + with network interface names as the keys and the namedtuple + described above as the values. + + If *nowrap* is True it detects and adjust the numbers which overflow + and wrap (restart from 0) and add "old value" to "new value" so that + the returned numbers will always be increasing or remain the same, + but never decrease. + "net_io_counters.cache_clear()" can be used to invalidate the + cache. + """ + rawdict = _psplatform.net_io_counters() + if not rawdict: + return {} if pernic else None + if nowrap: + rawdict = _wrap_numbers(rawdict, 'psutil.net_io_counters') + if pernic: + for nic, fields in rawdict.items(): + rawdict[nic] = _ntp.snetio(*fields) + return rawdict + else: + return _ntp.snetio(*[sum(x) for x in zip(*rawdict.values())]) + + +net_io_counters.cache_clear = functools.partial( + _wrap_numbers.cache_clear, 'psutil.net_io_counters' +) +net_io_counters.cache_clear.__doc__ = "Clears nowrap argument cache" + + +def net_connections(kind='inet'): + """Return system-wide socket connections as a list of + (fd, family, type, laddr, raddr, status, pid) namedtuples. + In case of limited privileges 'fd' and 'pid' may be set to -1 + and None respectively. + The *kind* parameter filters for connections that fit the + following criteria: + + +------------+----------------------------------------------------+ + | Kind Value | Connections using | + +------------+----------------------------------------------------+ + | inet | IPv4 and IPv6 | + | inet4 | IPv4 | + | inet6 | IPv6 | + | tcp | TCP | + | tcp4 | TCP over IPv4 | + | tcp6 | TCP over IPv6 | + | udp | UDP | + | udp4 | UDP over IPv4 | + | udp6 | UDP over IPv6 | + | unix | UNIX socket (both UDP and TCP protocols) | + | all | the sum of all the possible families and protocols | + +------------+----------------------------------------------------+ + + On macOS this function requires root privileges. + """ + _check_conn_kind(kind) + return _psplatform.net_connections(kind) + + +def net_if_addrs(): + """Return the addresses associated to each NIC (network interface + card) installed on the system as a dictionary whose keys are the + NIC names and value is a list of namedtuples for each address + assigned to the NIC. Each namedtuple includes 5 fields: + + - family: can be either socket.AF_INET, socket.AF_INET6 or + psutil.AF_LINK, which refers to a MAC address. + - address: is the primary address and it is always set. + - netmask: and 'broadcast' and 'ptp' may be None. + - ptp: stands for "point to point" and references the + destination address on a point to point interface + (typically a VPN). + - broadcast: and *ptp* are mutually exclusive. + + Note: you can have more than one address of the same family + associated with each interface. + """ + rawlist = _psplatform.net_if_addrs() + rawlist.sort(key=lambda x: x[1]) # sort by family + ret = collections.defaultdict(list) + for name, fam, addr, mask, broadcast, ptp in rawlist: + try: + fam = socket.AddressFamily(fam) + except ValueError: + if WINDOWS and fam == -1: + fam = _psplatform.AF_LINK + elif ( + hasattr(_psplatform, "AF_LINK") and fam == _psplatform.AF_LINK + ): + # Linux defines AF_LINK as an alias for AF_PACKET. + # We re-set the family here so that repr(family) + # will show AF_LINK rather than AF_PACKET + fam = _psplatform.AF_LINK + + if fam == _psplatform.AF_LINK: + # The underlying C function may return an incomplete MAC + # address in which case we fill it with null bytes, see: + # https://github.com/giampaolo/psutil/issues/786 + separator = ":" if POSIX else "-" + while addr.count(separator) < 5: + addr += f"{separator}00" + + nt = _ntp.snicaddr(fam, addr, mask, broadcast, ptp) + + # On Windows broadcast is None, so we determine it via + # ipaddress module. + if WINDOWS and fam in {socket.AF_INET, socket.AF_INET6}: + try: + broadcast = _common.broadcast_addr(nt) + except Exception as err: # noqa: BLE001 + debug(err) + else: + if broadcast is not None: + nt._replace(broadcast=broadcast) + + ret[name].append(nt) + + return dict(ret) + + +def net_if_stats(): + """Return information about each NIC (network interface card) + installed on the system as a dictionary whose keys are the + NIC names and value is a namedtuple with the following fields: + + - isup: whether the interface is up (bool) + - duplex: can be either NIC_DUPLEX_FULL, NIC_DUPLEX_HALF or + NIC_DUPLEX_UNKNOWN + - speed: the NIC speed expressed in mega bits (MB); if it can't + be determined (e.g. 'localhost') it will be set to 0. + - mtu: the maximum transmission unit expressed in bytes. + """ + return _psplatform.net_if_stats() + + +# ===================================================================== +# --- sensors +# ===================================================================== + + +# Linux, macOS +if hasattr(_psplatform, "sensors_temperatures"): + + def sensors_temperatures(fahrenheit=False): + """Return hardware temperatures. Each entry is a namedtuple + representing a certain hardware sensor (it may be a CPU, an + hard disk or something else, depending on the OS and its + configuration). + All temperatures are expressed in celsius unless *fahrenheit* + is set to True. + """ + + def convert(n): + if n is not None: + return (float(n) * 9 / 5) + 32 if fahrenheit else n + + ret = collections.defaultdict(list) + rawdict = _psplatform.sensors_temperatures() + + for name, values in rawdict.items(): + while values: + label, current, high, critical = values.pop(0) + current = convert(current) + high = convert(high) + critical = convert(critical) + + if high and not critical: + critical = high + elif critical and not high: + high = critical + + ret[name].append(_ntp.shwtemp(label, current, high, critical)) + + return dict(ret) + + __all__.append("sensors_temperatures") + + +# Linux +if hasattr(_psplatform, "sensors_fans"): + + def sensors_fans(): + """Return fans speed. Each entry is a namedtuple + representing a certain hardware sensor. + All speed are expressed in RPM (rounds per minute). + """ + return _psplatform.sensors_fans() + + __all__.append("sensors_fans") + + +# Linux, Windows, FreeBSD, macOS +if hasattr(_psplatform, "sensors_battery"): + + def sensors_battery(): + """Return battery information. If no battery is installed + returns None. + + - percent: battery power left as a percentage. + - secsleft: a rough approximation of how many seconds are left + before the battery runs out of power. May be + POWER_TIME_UNLIMITED or POWER_TIME_UNLIMITED. + - power_plugged: True if the AC power cable is connected. + """ + return _psplatform.sensors_battery() + + __all__.append("sensors_battery") + + +# ===================================================================== +# --- other system related functions +# ===================================================================== + + +def boot_time(): + """Return the system boot time expressed in seconds since the epoch + (seconds since January 1, 1970, at midnight UTC). The returned + value is based on the system clock, which means it may be affected + by changes such as manual adjustments or time synchronization (e.g. + NTP). + """ + return _psplatform.boot_time() + + +def users(): + """Return users currently connected on the system as a list of + namedtuples including the following fields. + + - user: the name of the user + - terminal: the tty or pseudo-tty associated with the user, if any. + - host: the host name associated with the entry, if any. + - started: the creation time as a floating point number expressed in + seconds since the epoch. + """ + return _psplatform.users() + + +# ===================================================================== +# --- Windows services +# ===================================================================== + + +if WINDOWS: + + def win_service_iter(): + """Return a generator yielding a WindowsService instance for all + Windows services installed. + """ + return _psplatform.win_service_iter() + + def win_service_get(name): + """Get a Windows service by *name*. + Raise NoSuchProcess if no service with such name exists. + """ + return _psplatform.win_service_get(name) + + +# ===================================================================== +# --- malloc / heap +# ===================================================================== + + +# Linux + glibc, Windows, macOS, FreeBSD, NetBSD +if hasattr(_psplatform, "heap_info"): + + def heap_info(): + """Return low-level heap statistics from the C heap allocator + (glibc). + + - `heap_used`: the total number of bytes allocated via + malloc/free. These are typically allocations smaller than + MMAP_THRESHOLD. + + - `mmap_used`: the total number of bytes allocated via `mmap()` + or via large ``malloc()`` allocations. + + - `heap_count` (Windows only): number of private heaps created + via `HeapCreate()`. + """ + return _ntp.pheap(*_psplatform.heap_info()) + + def heap_trim(): + """Request that the underlying allocator free any unused memory + it's holding in the heap (typically small `malloc()` + allocations). + + In practice, modern allocators rarely comply, so this is not a + general-purpose memory-reduction tool and won't meaningfully + shrink RSS in real programs. Its primary value is in **leak + detection tools**. + + Calling `heap_trim()` before taking measurements helps reduce + allocator noise, giving you a cleaner baseline so that changes + in `heap_used` come from the code you're testing, not from + internal allocator caching or fragmentation. Its effectiveness + depends on allocator behavior and fragmentation patterns. + """ + _psplatform.heap_trim() + + __all__.append("heap_info") + __all__.append("heap_trim") + + +# ===================================================================== + + +def _set_debug(value): + """Enable or disable PSUTIL_DEBUG option, which prints debugging + messages to stderr. + """ + import psutil._common + + psutil._common.PSUTIL_DEBUG = bool(value) + _psplatform.cext.set_debug(bool(value)) + + +del memoize_when_activated diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/psutil/_common.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/psutil/_common.py new file mode 100644 index 0000000000000000000000000000000000000000..33b6282253ac3321193fdd882aaf38185dee9d49 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/psutil/_common.py @@ -0,0 +1,861 @@ +# Copyright (c) 2009, Giampaolo Rodola'. All rights reserved. +# Use of this source code is governed by a BSD-style license that can be +# found in the LICENSE file. + +"""Common objects shared by __init__.py and _ps*.py modules. + +Note: this module is imported by setup.py, so it should not import +psutil or third-party modules. +""" + +import collections +import enum +import functools +import os +import socket +import stat +import sys +import threading +import warnings +from socket import AF_INET +from socket import SOCK_DGRAM +from socket import SOCK_STREAM + +try: + from socket import AF_INET6 +except ImportError: + AF_INET6 = None +try: + from socket import AF_UNIX +except ImportError: + AF_UNIX = None + + +PSUTIL_DEBUG = bool(os.getenv('PSUTIL_DEBUG')) +_DEFAULT = object() + +# fmt: off +__all__ = [ + # OS constants + 'FREEBSD', 'BSD', 'LINUX', 'NETBSD', 'OPENBSD', 'MACOS', 'OSX', 'POSIX', + 'SUNOS', 'WINDOWS', + # connection constants + 'CONN_CLOSE', 'CONN_CLOSE_WAIT', 'CONN_CLOSING', 'CONN_ESTABLISHED', + 'CONN_FIN_WAIT1', 'CONN_FIN_WAIT2', 'CONN_LAST_ACK', 'CONN_LISTEN', + 'CONN_NONE', 'CONN_SYN_RECV', 'CONN_SYN_SENT', 'CONN_TIME_WAIT', + # net constants + 'NIC_DUPLEX_FULL', 'NIC_DUPLEX_HALF', 'NIC_DUPLEX_UNKNOWN', # noqa: F822 + # process status constants + 'STATUS_DEAD', 'STATUS_DISK_SLEEP', 'STATUS_IDLE', 'STATUS_LOCKED', + 'STATUS_RUNNING', 'STATUS_SLEEPING', 'STATUS_STOPPED', 'STATUS_SUSPENDED', + 'STATUS_TRACING_STOP', 'STATUS_WAITING', 'STATUS_WAKE_KILL', + 'STATUS_WAKING', 'STATUS_ZOMBIE', 'STATUS_PARKED', + # other constants + 'ENCODING', 'ENCODING_ERRS', 'AF_INET6', + # utility functions + 'conn_tmap', 'deprecated_method', 'isfile_strict', 'memoize', + 'parse_environ_block', 'path_exists_strict', 'usage_percent', + 'supports_ipv6', 'sockfam_to_enum', 'socktype_to_enum', "wrap_numbers", + 'open_text', 'open_binary', 'cat', 'bcat', + 'bytes2human', 'conn_to_ntuple', 'debug', + # shell utils + 'hilite', 'term_supports_colors', 'print_color', +] +# fmt: on + + +# =================================================================== +# --- OS constants +# =================================================================== + + +POSIX = os.name == "posix" +WINDOWS = os.name == "nt" +LINUX = sys.platform.startswith("linux") +MACOS = sys.platform.startswith("darwin") +OSX = MACOS # deprecated alias +FREEBSD = sys.platform.startswith(("freebsd", "midnightbsd")) +OPENBSD = sys.platform.startswith("openbsd") +NETBSD = sys.platform.startswith("netbsd") +BSD = FREEBSD or OPENBSD or NETBSD +SUNOS = sys.platform.startswith(("sunos", "solaris")) +AIX = sys.platform.startswith("aix") + + +# =================================================================== +# --- API constants +# =================================================================== + + +# Process.status() +STATUS_RUNNING = "running" +STATUS_SLEEPING = "sleeping" +STATUS_DISK_SLEEP = "disk-sleep" +STATUS_STOPPED = "stopped" +STATUS_TRACING_STOP = "tracing-stop" +STATUS_ZOMBIE = "zombie" +STATUS_DEAD = "dead" +STATUS_WAKE_KILL = "wake-kill" +STATUS_WAKING = "waking" +STATUS_IDLE = "idle" # Linux, macOS, FreeBSD +STATUS_LOCKED = "locked" # FreeBSD +STATUS_WAITING = "waiting" # FreeBSD +STATUS_SUSPENDED = "suspended" # NetBSD +STATUS_PARKED = "parked" # Linux + +# Process.net_connections() and psutil.net_connections() +CONN_ESTABLISHED = "ESTABLISHED" +CONN_SYN_SENT = "SYN_SENT" +CONN_SYN_RECV = "SYN_RECV" +CONN_FIN_WAIT1 = "FIN_WAIT1" +CONN_FIN_WAIT2 = "FIN_WAIT2" +CONN_TIME_WAIT = "TIME_WAIT" +CONN_CLOSE = "CLOSE" +CONN_CLOSE_WAIT = "CLOSE_WAIT" +CONN_LAST_ACK = "LAST_ACK" +CONN_LISTEN = "LISTEN" +CONN_CLOSING = "CLOSING" +CONN_NONE = "NONE" + + +# net_if_stats() +class NicDuplex(enum.IntEnum): + NIC_DUPLEX_FULL = 2 + NIC_DUPLEX_HALF = 1 + NIC_DUPLEX_UNKNOWN = 0 + + +globals().update(NicDuplex.__members__) + + +# sensors_battery() +class BatteryTime(enum.IntEnum): + POWER_TIME_UNKNOWN = -1 + POWER_TIME_UNLIMITED = -2 + + +globals().update(BatteryTime.__members__) + +# --- others + +ENCODING = sys.getfilesystemencoding() +ENCODING_ERRS = sys.getfilesystemencodeerrors() + + +# =================================================================== +# --- Process.net_connections() 'kind' parameter mapping +# =================================================================== + + +conn_tmap = { + "all": ([AF_INET, AF_INET6, AF_UNIX], [SOCK_STREAM, SOCK_DGRAM]), + "tcp": ([AF_INET, AF_INET6], [SOCK_STREAM]), + "tcp4": ([AF_INET], [SOCK_STREAM]), + "udp": ([AF_INET, AF_INET6], [SOCK_DGRAM]), + "udp4": ([AF_INET], [SOCK_DGRAM]), + "inet": ([AF_INET, AF_INET6], [SOCK_STREAM, SOCK_DGRAM]), + "inet4": ([AF_INET], [SOCK_STREAM, SOCK_DGRAM]), + "inet6": ([AF_INET6], [SOCK_STREAM, SOCK_DGRAM]), +} + +if AF_INET6 is not None: + conn_tmap.update({ + "tcp6": ([AF_INET6], [SOCK_STREAM]), + "udp6": ([AF_INET6], [SOCK_DGRAM]), + }) + +if AF_UNIX is not None and not SUNOS: + conn_tmap.update({"unix": ([AF_UNIX], [SOCK_STREAM, SOCK_DGRAM])}) + + +# ===================================================================== +# --- Exceptions +# ===================================================================== + + +class Error(Exception): + """Base exception class. All other psutil exceptions inherit + from this one. + """ + + __module__ = 'psutil' + + def _infodict(self, attrs): + info = collections.OrderedDict() + for name in attrs: + value = getattr(self, name, None) + if value or (name == "pid" and value == 0): + info[name] = value + return info + + def __str__(self): + # invoked on `raise Error` + info = self._infodict(("pid", "ppid", "name")) + if info: + details = "({})".format( + ", ".join([f"{k}={v!r}" for k, v in info.items()]) + ) + else: + details = None + return " ".join([x for x in (getattr(self, "msg", ""), details) if x]) + + def __repr__(self): + # invoked on `repr(Error)` + info = self._infodict(("pid", "ppid", "name", "seconds", "msg")) + details = ", ".join([f"{k}={v!r}" for k, v in info.items()]) + return f"psutil.{self.__class__.__name__}({details})" + + +class NoSuchProcess(Error): + """Exception raised when a process with a certain PID doesn't + or no longer exists. + """ + + __module__ = 'psutil' + + def __init__(self, pid, name=None, msg=None): + Error.__init__(self) + self.pid = pid + self.name = name + self.msg = msg or "process no longer exists" + + def __reduce__(self): + return (self.__class__, (self.pid, self.name, self.msg)) + + +class ZombieProcess(NoSuchProcess): + """Exception raised when querying a zombie process. This is + raised on macOS, BSD and Solaris only, and not always: depending + on the query the OS may be able to succeed anyway. + On Linux all zombie processes are querable (hence this is never + raised). Windows doesn't have zombie processes. + """ + + __module__ = 'psutil' + + def __init__(self, pid, name=None, ppid=None, msg=None): + NoSuchProcess.__init__(self, pid, name, msg) + self.ppid = ppid + self.msg = msg or "PID still exists but it's a zombie" + + def __reduce__(self): + return (self.__class__, (self.pid, self.name, self.ppid, self.msg)) + + +class AccessDenied(Error): + """Exception raised when permission to perform an action is denied.""" + + __module__ = 'psutil' + + def __init__(self, pid=None, name=None, msg=None): + Error.__init__(self) + self.pid = pid + self.name = name + self.msg = msg or "" + + def __reduce__(self): + return (self.__class__, (self.pid, self.name, self.msg)) + + +class TimeoutExpired(Error): + """Raised on Process.wait(timeout) if timeout expires and process + is still alive. + """ + + __module__ = 'psutil' + + def __init__(self, seconds, pid=None, name=None): + Error.__init__(self) + self.seconds = seconds + self.pid = pid + self.name = name + self.msg = f"timeout after {seconds} seconds" + + def __reduce__(self): + return (self.__class__, (self.seconds, self.pid, self.name)) + + +# =================================================================== +# --- utils +# =================================================================== + + +def usage_percent(used, total, round_=None): + """Calculate percentage usage of 'used' against 'total'.""" + try: + ret = (float(used) / total) * 100 + except ZeroDivisionError: + return 0.0 + else: + if round_ is not None: + ret = round(ret, round_) + return ret + + +def memoize(fun): + """A simple memoize decorator for functions supporting (hashable) + positional arguments. + It also provides a cache_clear() function for clearing the cache: + + >>> @memoize + ... def foo() + ... return 1 + ... + >>> foo() + 1 + >>> foo.cache_clear() + >>> + + It supports: + - functions + - classes (acts as a @singleton) + - staticmethods + - classmethods + + It does NOT support: + - methods + """ + + @functools.wraps(fun) + def wrapper(*args, **kwargs): + key = (args, frozenset(sorted(kwargs.items()))) + try: + return cache[key] + except KeyError: + try: + ret = cache[key] = fun(*args, **kwargs) + except Exception as err: + raise err from None + return ret + + def cache_clear(): + """Clear cache.""" + cache.clear() + + cache = {} + wrapper.cache_clear = cache_clear + return wrapper + + +def memoize_when_activated(fun): + """A memoize decorator which is disabled by default. It can be + activated and deactivated on request. + For efficiency reasons it can be used only against class methods + accepting no arguments. + + >>> class Foo: + ... @memoize + ... def foo() + ... print(1) + ... + >>> f = Foo() + >>> # deactivated (default) + >>> foo() + 1 + >>> foo() + 1 + >>> + >>> # activated + >>> foo.cache_activate(self) + >>> foo() + 1 + >>> foo() + >>> foo() + >>> + """ + + @functools.wraps(fun) + def wrapper(self): + try: + # case 1: we previously entered oneshot() ctx + ret = self._cache[fun] + except AttributeError: + # case 2: we never entered oneshot() ctx + try: + return fun(self) + except Exception as err: + raise err from None + except KeyError: + # case 3: we entered oneshot() ctx but there's no cache + # for this entry yet + try: + ret = fun(self) + except Exception as err: + raise err from None + try: + self._cache[fun] = ret + except AttributeError: + # multi-threading race condition, see: + # https://github.com/giampaolo/psutil/issues/1948 + pass + return ret + + def cache_activate(proc): + """Activate cache. Expects a Process instance. Cache will be + stored as a "_cache" instance attribute. + """ + proc._cache = {} + + def cache_deactivate(proc): + """Deactivate and clear cache.""" + try: + del proc._cache + except AttributeError: + pass + + wrapper.cache_activate = cache_activate + wrapper.cache_deactivate = cache_deactivate + return wrapper + + +def isfile_strict(path): + """Same as os.path.isfile() but does not swallow EACCES / EPERM + exceptions, see: + http://mail.python.org/pipermail/python-dev/2012-June/120787.html. + """ + try: + st = os.stat(path) + except PermissionError: + raise + except OSError: + return False + else: + return stat.S_ISREG(st.st_mode) + + +def path_exists_strict(path): + """Same as os.path.exists() but does not swallow EACCES / EPERM + exceptions. See: + http://mail.python.org/pipermail/python-dev/2012-June/120787.html. + """ + try: + os.stat(path) + except PermissionError: + raise + except OSError: + return False + else: + return True + + +def supports_ipv6(): + """Return True if IPv6 is supported on this platform.""" + if not socket.has_ipv6 or AF_INET6 is None: + return False + try: + with socket.socket(AF_INET6, socket.SOCK_STREAM) as sock: + sock.bind(("::1", 0)) + return True + except OSError: + return False + + +def parse_environ_block(data): + """Parse a C environ block of environment variables into a dictionary.""" + # The block is usually raw data from the target process. It might contain + # trailing garbage and lines that do not look like assignments. + ret = {} + pos = 0 + + # localize global variable to speed up access. + WINDOWS_ = WINDOWS + while True: + next_pos = data.find("\0", pos) + # nul byte at the beginning or double nul byte means finish + if next_pos <= pos: + break + # there might not be an equals sign + equal_pos = data.find("=", pos, next_pos) + if equal_pos > pos: + key = data[pos:equal_pos] + value = data[equal_pos + 1 : next_pos] + # Windows expects environment variables to be uppercase only + if WINDOWS_: + key = key.upper() + ret[key] = value + pos = next_pos + 1 + + return ret + + +def sockfam_to_enum(num): + """Convert a numeric socket family value to an IntEnum member. + If it's not a known member, return the numeric value itself. + """ + try: + return socket.AddressFamily(num) + except ValueError: + return num + + +def socktype_to_enum(num): + """Convert a numeric socket type value to an IntEnum member. + If it's not a known member, return the numeric value itself. + """ + try: + return socket.SocketKind(num) + except ValueError: + return num + + +def conn_to_ntuple(fd, fam, type_, laddr, raddr, status, status_map, pid=None): + """Convert a raw connection tuple to a proper ntuple.""" + from . import _ntuples as ntp + + if fam in {socket.AF_INET, AF_INET6}: + if laddr: + laddr = ntp.addr(*laddr) + if raddr: + raddr = ntp.addr(*raddr) + if type_ == socket.SOCK_STREAM and fam in {AF_INET, AF_INET6}: + status = status_map.get(status, CONN_NONE) + else: + status = CONN_NONE # ignore whatever C returned to us + fam = sockfam_to_enum(fam) + type_ = socktype_to_enum(type_) + if pid is None: + return ntp.pconn(fd, fam, type_, laddr, raddr, status) + else: + return ntp.sconn(fd, fam, type_, laddr, raddr, status, pid) + + +def broadcast_addr(addr): + """Given the address ntuple returned by ``net_if_addrs()`` + calculates the broadcast address. + """ + import ipaddress + + if not addr.address or not addr.netmask: + return None + if addr.family == socket.AF_INET: + return str( + ipaddress.IPv4Network( + f"{addr.address}/{addr.netmask}", strict=False + ).broadcast_address + ) + if addr.family == socket.AF_INET6: + return str( + ipaddress.IPv6Network( + f"{addr.address}/{addr.netmask}", strict=False + ).broadcast_address + ) + + +def deprecated_method(replacement): + """A decorator which can be used to mark a method as deprecated + 'replcement' is the method name which will be called instead. + """ + + def outer(fun): + msg = ( + f"{fun.__name__}() is deprecated and will be removed; use" + f" {replacement}() instead" + ) + if fun.__doc__ is None: + fun.__doc__ = msg + + @functools.wraps(fun) + def inner(self, *args, **kwargs): + warnings.warn(msg, category=DeprecationWarning, stacklevel=2) + return getattr(self, replacement)(*args, **kwargs) + + return inner + + return outer + + +class _WrapNumbers: + """Watches numbers so that they don't overflow and wrap + (reset to zero). + """ + + def __init__(self): + self.lock = threading.Lock() + self.cache = {} + self.reminders = {} + self.reminder_keys = {} + + def _add_dict(self, input_dict, name): + assert name not in self.cache + assert name not in self.reminders + assert name not in self.reminder_keys + self.cache[name] = input_dict + self.reminders[name] = collections.defaultdict(int) + self.reminder_keys[name] = collections.defaultdict(set) + + def _remove_dead_reminders(self, input_dict, name): + """In case the number of keys changed between calls (e.g. a + disk disappears) this removes the entry from self.reminders. + """ + old_dict = self.cache[name] + gone_keys = set(old_dict.keys()) - set(input_dict.keys()) + for gone_key in gone_keys: + for remkey in self.reminder_keys[name][gone_key]: + del self.reminders[name][remkey] + del self.reminder_keys[name][gone_key] + + def run(self, input_dict, name): + """Cache dict and sum numbers which overflow and wrap. + Return an updated copy of `input_dict`. + """ + if name not in self.cache: + # This was the first call. + self._add_dict(input_dict, name) + return input_dict + + self._remove_dead_reminders(input_dict, name) + + old_dict = self.cache[name] + new_dict = {} + for key in input_dict: + input_tuple = input_dict[key] + try: + old_tuple = old_dict[key] + except KeyError: + # The input dict has a new key (e.g. a new disk or NIC) + # which didn't exist in the previous call. + new_dict[key] = input_tuple + continue + + bits = [] + for i in range(len(input_tuple)): + input_value = input_tuple[i] + old_value = old_tuple[i] + remkey = (key, i) + if input_value < old_value: + # it wrapped! + self.reminders[name][remkey] += old_value + self.reminder_keys[name][key].add(remkey) + bits.append(input_value + self.reminders[name][remkey]) + + new_dict[key] = tuple(bits) + + self.cache[name] = input_dict + return new_dict + + def cache_clear(self, name=None): + """Clear the internal cache, optionally only for function 'name'.""" + with self.lock: + if name is None: + self.cache.clear() + self.reminders.clear() + self.reminder_keys.clear() + else: + self.cache.pop(name, None) + self.reminders.pop(name, None) + self.reminder_keys.pop(name, None) + + def cache_info(self): + """Return internal cache dicts as a tuple of 3 elements.""" + with self.lock: + return (self.cache, self.reminders, self.reminder_keys) + + +def wrap_numbers(input_dict, name): + """Given an `input_dict` and a function `name`, adjust the numbers + which "wrap" (restart from zero) across different calls by adding + "old value" to "new value" and return an updated dict. + """ + with _wn.lock: + return _wn.run(input_dict, name) + + +_wn = _WrapNumbers() +wrap_numbers.cache_clear = _wn.cache_clear +wrap_numbers.cache_info = _wn.cache_info + + +# The read buffer size for open() builtin. This (also) dictates how +# much data we read(2) when iterating over file lines as in: +# >>> with open(file) as f: +# ... for line in f: +# ... ... +# Default per-line buffer size for binary files is 1K. For text files +# is 8K. We use a bigger buffer (32K) in order to have more consistent +# results when reading /proc pseudo files on Linux, see: +# https://github.com/giampaolo/psutil/issues/2050 +# https://github.com/giampaolo/psutil/issues/708 +FILE_READ_BUFFER_SIZE = 32 * 1024 + + +def open_binary(fname): + return open(fname, "rb", buffering=FILE_READ_BUFFER_SIZE) + + +def open_text(fname): + """Open a file in text mode by using the proper FS encoding and + en/decoding error handlers. + """ + # See: + # https://github.com/giampaolo/psutil/issues/675 + # https://github.com/giampaolo/psutil/pull/733 + fobj = open( # noqa: SIM115 + fname, + buffering=FILE_READ_BUFFER_SIZE, + encoding=ENCODING, + errors=ENCODING_ERRS, + ) + try: + # Dictates per-line read(2) buffer size. Defaults is 8k. See: + # https://github.com/giampaolo/psutil/issues/2050#issuecomment-1013387546 + fobj._CHUNK_SIZE = FILE_READ_BUFFER_SIZE + except AttributeError: + pass + except Exception: + fobj.close() + raise + + return fobj + + +def cat(fname, fallback=_DEFAULT, _open=open_text): + """Read entire file content and return it as a string. File is + opened in text mode. If specified, `fallback` is the value + returned in case of error, either if the file does not exist or + it can't be read(). + """ + if fallback is _DEFAULT: + with _open(fname) as f: + return f.read() + else: + try: + with _open(fname) as f: + return f.read() + except OSError: + return fallback + + +def bcat(fname, fallback=_DEFAULT): + """Same as above but opens file in binary mode.""" + return cat(fname, fallback=fallback, _open=open_binary) + + +def bytes2human(n, format="%(value).1f%(symbol)s"): + """Used by various scripts. See: https://code.activestate.com/recipes/578019-bytes-to-human-human-to-bytes-converter/?in=user-4178764. + + >>> bytes2human(10000) + '9.8K' + >>> bytes2human(100001221) + '95.4M' + """ + symbols = ('B', 'K', 'M', 'G', 'T', 'P', 'E', 'Z', 'Y') + prefix = {} + for i, s in enumerate(symbols[1:]): + prefix[s] = 1 << (i + 1) * 10 + for symbol in reversed(symbols[1:]): + if abs(n) >= prefix[symbol]: + value = float(n) / prefix[symbol] + return format % locals() + return format % dict(symbol=symbols[0], value=n) + + +def get_procfs_path(): + """Return updated psutil.PROCFS_PATH constant.""" + return sys.modules['psutil'].PROCFS_PATH + + +def decode(s): + return s.decode(encoding=ENCODING, errors=ENCODING_ERRS) + + +# ===================================================================== +# --- shell utils +# ===================================================================== + + +@memoize +def term_supports_colors(file=sys.stdout): # pragma: no cover + if not hasattr(file, "isatty") or not file.isatty(): + return False + try: + file.fileno() + except Exception: # noqa: BLE001 + return False + return True + + +def hilite(s, color=None, bold=False): # pragma: no cover + """Return an highlighted version of 'string'.""" + if not term_supports_colors(): + return s + attr = [] + colors = dict( + blue='34', + brown='33', + darkgrey='30', + green='32', + grey='37', + lightblue='36', + red='91', + violet='35', + yellow='93', + ) + colors[None] = '29' + try: + color = colors[color] + except KeyError: + msg = f"invalid color {color!r}; choose amongst {list(colors.keys())}" + raise ValueError(msg) from None + attr.append(color) + if bold: + attr.append('1') + return f"\x1b[{';'.join(attr)}m{s}\x1b[0m" + + +def print_color( + s, color=None, bold=False, file=sys.stdout +): # pragma: no cover + """Print a colorized version of string.""" + if not term_supports_colors(): + print(s, file=file) + elif POSIX: + print(hilite(s, color, bold), file=file) + else: + import ctypes + + DEFAULT_COLOR = 7 + GetStdHandle = ctypes.windll.Kernel32.GetStdHandle + SetConsoleTextAttribute = ( + ctypes.windll.Kernel32.SetConsoleTextAttribute + ) + + colors = dict(green=2, red=4, brown=6, yellow=6) + colors[None] = DEFAULT_COLOR + try: + color = colors[color] + except KeyError: + msg = ( + f"invalid color {color!r}; choose between" + f" {list(colors.keys())!r}" + ) + raise ValueError(msg) from None + if bold and color <= 7: + color += 8 + + handle_id = -12 if file is sys.stderr else -11 + GetStdHandle.restype = ctypes.c_ulong + handle = GetStdHandle(handle_id) + SetConsoleTextAttribute(handle, color) + try: + print(s, file=file) + finally: + SetConsoleTextAttribute(handle, DEFAULT_COLOR) + + +def debug(msg): + """If PSUTIL_DEBUG env var is set, print a debug message to stderr.""" + if PSUTIL_DEBUG: + import inspect + + fname, lineno, _, _lines, _index = inspect.getframeinfo( + inspect.currentframe().f_back + ) + if isinstance(msg, Exception): + if isinstance(msg, OSError): + # ...because str(exc) may contain info about the file name + msg = f"ignoring {msg}" + else: + msg = f"ignoring {msg!r}" + print( # noqa: T201 + f"psutil-debug [{fname}:{lineno}]> {msg}", file=sys.stderr + ) diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/psutil/_ntuples.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/psutil/_ntuples.py new file mode 100644 index 0000000000000000000000000000000000000000..f0ab0393763f0969d2f4b4302fa638771c6c0bdb --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/psutil/_ntuples.py @@ -0,0 +1,429 @@ +# Copyright (c) 2009, Giampaolo Rodola". All rights reserved. +# Use of this source code is governed by a BSD-style license that can be +# found in the LICENSE file. + +from collections import namedtuple as nt + +from ._common import AIX +from ._common import BSD +from ._common import FREEBSD +from ._common import LINUX +from ._common import MACOS +from ._common import SUNOS +from ._common import WINDOWS + +# =================================================================== +# --- system functions +# =================================================================== + +# psutil.swap_memory() +sswap = nt("sswap", ("total", "used", "free", "percent", "sin", "sout")) + +# psutil.disk_usage() +sdiskusage = nt("sdiskusage", ("total", "used", "free", "percent")) + +# psutil.disk_io_counters() +sdiskio = nt( + "sdiskio", + ( + "read_count", + "write_count", + "read_bytes", + "write_bytes", + "read_time", + "write_time", + ), +) + +# psutil.disk_partitions() +sdiskpart = nt("sdiskpart", ("device", "mountpoint", "fstype", "opts")) + +# psutil.net_io_counters() +snetio = nt( + "snetio", + ( + "bytes_sent", + "bytes_recv", + "packets_sent", + "packets_recv", + "errin", + "errout", + "dropin", + "dropout", + ), +) + +# psutil.users() +suser = nt("suser", ("name", "terminal", "host", "started", "pid")) + +# psutil.net_connections() +sconn = nt( + "sconn", ("fd", "family", "type", "laddr", "raddr", "status", "pid") +) + +# psutil.net_if_addrs() +snicaddr = nt("snicaddr", ("family", "address", "netmask", "broadcast", "ptp")) + +# psutil.net_if_stats() +snicstats = nt("snicstats", ("isup", "duplex", "speed", "mtu", "flags")) + +# psutil.cpu_stats() +scpustats = nt( + "scpustats", ("ctx_switches", "interrupts", "soft_interrupts", "syscalls") +) + +# psutil.cpu_freq() +scpufreq = nt("scpufreq", ("current", "min", "max")) + +# psutil.sensors_temperatures() +shwtemp = nt("shwtemp", ("label", "current", "high", "critical")) + +# psutil.sensors_battery() +sbattery = nt("sbattery", ("percent", "secsleft", "power_plugged")) + +# psutil.sensors_fans() +sfan = nt("sfan", ("label", "current")) + +# psutil.heap_info() (mallinfo2 Linux struct) +if LINUX or WINDOWS or MACOS or BSD: + pheap = nt( + "pheap", + [ + "heap_used", # uordblks, memory allocated via malloc() + "mmap_used", # hblkhd, memory allocated via mmap() (large blocks) + ], + ) + if WINDOWS: + pheap = nt("pheap", pheap._fields + ("heap_count",)) + +# =================================================================== +# --- Process class +# =================================================================== + +# psutil.Process.cpu_times() +pcputimes = nt( + "pcputimes", ("user", "system", "children_user", "children_system") +) + +# psutil.Process.open_files() +popenfile = nt("popenfile", ("path", "fd")) + +# psutil.Process.threads() +pthread = nt("pthread", ("id", "user_time", "system_time")) + +# psutil.Process.uids() +puids = nt("puids", ("real", "effective", "saved")) + +# psutil.Process.gids() +pgids = nt("pgids", ("real", "effective", "saved")) + +# psutil.Process.io_counters() +pio = nt("pio", ("read_count", "write_count", "read_bytes", "write_bytes")) + +# psutil.Process.ionice() +pionice = nt("pionice", ("ioclass", "value")) + +# psutil.Process.ctx_switches() +pctxsw = nt("pctxsw", ("voluntary", "involuntary")) + +# psutil.Process.net_connections() +pconn = nt("pconn", ("fd", "family", "type", "laddr", "raddr", "status")) + +# psutil.net_connections() and psutil.Process.net_connections() +addr = nt("addr", ("ip", "port")) + +# =================================================================== +# --- Linux +# =================================================================== + +if LINUX: + + # This gets set from _pslinux.py + scputimes = None + + # psutil.virtual_memory() + svmem = nt( + "svmem", + ( + "total", + "available", + "percent", + "used", + "free", + "active", + "inactive", + "buffers", + "cached", + "shared", + "slab", + ), + ) + + # psutil.disk_io_counters() + sdiskio = nt( + "sdiskio", + ( + "read_count", + "write_count", + "read_bytes", + "write_bytes", + "read_time", + "write_time", + "read_merged_count", + "write_merged_count", + "busy_time", + ), + ) + + # psutil.Process().open_files() + popenfile = nt("popenfile", ("path", "fd", "position", "mode", "flags")) + + # psutil.Process().memory_info() + pmem = nt("pmem", ("rss", "vms", "shared", "text", "lib", "data", "dirty")) + + # psutil.Process().memory_full_info() + pfullmem = nt("pfullmem", pmem._fields + ("uss", "pss", "swap")) + + # psutil.Process().memory_maps(grouped=True) + pmmap_grouped = nt( + "pmmap_grouped", + ( + "path", + "rss", + "size", + "pss", + "shared_clean", + "shared_dirty", + "private_clean", + "private_dirty", + "referenced", + "anonymous", + "swap", + ), + ) + + # psutil.Process().memory_maps(grouped=False) + pmmap_ext = nt( + "pmmap_ext", "addr perms " + " ".join(pmmap_grouped._fields) + ) + + # psutil.Process.io_counters() + pio = nt( + "pio", + ( + "read_count", + "write_count", + "read_bytes", + "write_bytes", + "read_chars", + "write_chars", + ), + ) + + # psutil.Process.cpu_times() + pcputimes = nt( + "pcputimes", + ("user", "system", "children_user", "children_system", "iowait"), + ) + +# =================================================================== +# --- Windows +# =================================================================== + +elif WINDOWS: + + # psutil.cpu_times() + scputimes = nt("scputimes", ("user", "system", "idle", "interrupt", "dpc")) + + # psutil.virtual_memory() + svmem = nt("svmem", ("total", "available", "percent", "used", "free")) + + # psutil.Process.memory_info() + pmem = nt( + "pmem", + ( + "rss", + "vms", + "num_page_faults", + "peak_wset", + "wset", + "peak_paged_pool", + "paged_pool", + "peak_nonpaged_pool", + "nonpaged_pool", + "pagefile", + "peak_pagefile", + "private", + ), + ) + + # psutil.Process.memory_full_info() + pfullmem = nt("pfullmem", pmem._fields + ("uss",)) + + # psutil.Process.memory_maps(grouped=True) + pmmap_grouped = nt("pmmap_grouped", ("path", "rss")) + + # psutil.Process.memory_maps(grouped=False) + pmmap_ext = nt( + "pmmap_ext", "addr perms " + " ".join(pmmap_grouped._fields) + ) + + # psutil.Process.io_counters() + pio = nt( + "pio", + ( + "read_count", + "write_count", + "read_bytes", + "write_bytes", + "other_count", + "other_bytes", + ), + ) + +# =================================================================== +# --- macOS +# =================================================================== + +elif MACOS: + + # psutil.cpu_times() + scputimes = nt("scputimes", ("user", "nice", "system", "idle")) + + # psutil.virtual_memory() + svmem = nt( + "svmem", + ( + "total", + "available", + "percent", + "used", + "free", + "active", + "inactive", + "wired", + ), + ) + + # psutil.Process.memory_info() + pmem = nt("pmem", ("rss", "vms", "pfaults", "pageins")) + + # psutil.Process.memory_full_info() + pfullmem = nt("pfullmem", pmem._fields + ("uss",)) + +# =================================================================== +# --- BSD +# =================================================================== + +elif BSD: + + # psutil.virtual_memory() + svmem = nt( + "svmem", + ( + "total", + "available", + "percent", + "used", + "free", + "active", + "inactive", + "buffers", + "cached", + "shared", + "wired", + ), + ) + + # psutil.cpu_times() + scputimes = nt("scputimes", ("user", "nice", "system", "idle", "irq")) + + # psutil.Process.memory_info() + pmem = nt("pmem", ("rss", "vms", "text", "data", "stack")) + + # psutil.Process.memory_full_info() + pfullmem = pmem + + # psutil.Process.cpu_times() + pcputimes = nt( + "pcputimes", ("user", "system", "children_user", "children_system") + ) + + # psutil.Process.memory_maps(grouped=True) + pmmap_grouped = nt( + "pmmap_grouped", "path rss, private, ref_count, shadow_count" + ) + + # psutil.Process.memory_maps(grouped=False) + pmmap_ext = nt( + "pmmap_ext", "addr, perms path rss, private, ref_count, shadow_count" + ) + + # psutil.disk_io_counters() + if FREEBSD: + sdiskio = nt( + "sdiskio", + ( + "read_count", + "write_count", + "read_bytes", + "write_bytes", + "read_time", + "write_time", + "busy_time", + ), + ) + else: + sdiskio = nt( + "sdiskio", + ("read_count", "write_count", "read_bytes", "write_bytes"), + ) + +# =================================================================== +# --- SunOS +# =================================================================== + +elif SUNOS: + + # psutil.cpu_times() + scputimes = nt("scputimes", ("user", "system", "idle", "iowait")) + + # psutil.cpu_times(percpu=True) + pcputimes = nt( + "pcputimes", ("user", "system", "children_user", "children_system") + ) + + # psutil.virtual_memory() + svmem = nt("svmem", ("total", "available", "percent", "used", "free")) + + # psutil.Process.memory_info() + pmem = nt("pmem", ("rss", "vms")) + + # psutil.Process.memory_full_info() + pfullmem = pmem + + # psutil.Process.memory_maps(grouped=True) + pmmap_grouped = nt("pmmap_grouped", ("path", "rss", "anonymous", "locked")) + + # psutil.Process.memory_maps(grouped=False) + pmmap_ext = nt( + "pmmap_ext", "addr perms " + " ".join(pmmap_grouped._fields) + ) + +# =================================================================== +# --- AIX +# =================================================================== + +elif AIX: + + # psutil.Process.memory_info() + pmem = nt("pmem", ("rss", "vms")) + + # psutil.Process.memory_full_info() + pfullmem = pmem + + # psutil.Process.cpu_times() + scputimes = nt("scputimes", ("user", "system", "idle", "iowait")) + + # psutil.virtual_memory() + svmem = nt("svmem", ("total", "available", "percent", "used", "free")) diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/psutil/_psaix.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/psutil/_psaix.py new file mode 100644 index 0000000000000000000000000000000000000000..1f628ad93c1b1464b16d63daee79fe161204730e --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/psutil/_psaix.py @@ -0,0 +1,546 @@ +# Copyright (c) 2009, Giampaolo Rodola' +# Copyright (c) 2017, Arnon Yaari +# All rights reserved. +# Use of this source code is governed by a BSD-style license that can be +# found in the LICENSE file. + +"""AIX platform implementation.""" + +import functools +import glob +import os +import re +import subprocess +import sys + +from . import _common +from . import _ntuples as ntp +from . import _psposix +from . import _psutil_aix as cext +from ._common import NIC_DUPLEX_FULL +from ._common import NIC_DUPLEX_HALF +from ._common import NIC_DUPLEX_UNKNOWN +from ._common import AccessDenied +from ._common import NoSuchProcess +from ._common import ZombieProcess +from ._common import conn_to_ntuple +from ._common import get_procfs_path +from ._common import memoize_when_activated +from ._common import usage_percent + +__extra__all__ = ["PROCFS_PATH"] + + +# ===================================================================== +# --- globals +# ===================================================================== + + +HAS_THREADS = hasattr(cext, "proc_threads") +HAS_NET_IO_COUNTERS = hasattr(cext, "net_io_counters") +HAS_PROC_IO_COUNTERS = hasattr(cext, "proc_io_counters") + +PAGE_SIZE = cext.getpagesize() +AF_LINK = cext.AF_LINK + +PROC_STATUSES = { + cext.SIDL: _common.STATUS_IDLE, + cext.SZOMB: _common.STATUS_ZOMBIE, + cext.SACTIVE: _common.STATUS_RUNNING, + cext.SSWAP: _common.STATUS_RUNNING, # TODO what status is this? + cext.SSTOP: _common.STATUS_STOPPED, +} + +TCP_STATUSES = { + cext.TCPS_ESTABLISHED: _common.CONN_ESTABLISHED, + cext.TCPS_SYN_SENT: _common.CONN_SYN_SENT, + cext.TCPS_SYN_RCVD: _common.CONN_SYN_RECV, + cext.TCPS_FIN_WAIT_1: _common.CONN_FIN_WAIT1, + cext.TCPS_FIN_WAIT_2: _common.CONN_FIN_WAIT2, + cext.TCPS_TIME_WAIT: _common.CONN_TIME_WAIT, + cext.TCPS_CLOSED: _common.CONN_CLOSE, + cext.TCPS_CLOSE_WAIT: _common.CONN_CLOSE_WAIT, + cext.TCPS_LAST_ACK: _common.CONN_LAST_ACK, + cext.TCPS_LISTEN: _common.CONN_LISTEN, + cext.TCPS_CLOSING: _common.CONN_CLOSING, + cext.PSUTIL_CONN_NONE: _common.CONN_NONE, +} + +proc_info_map = dict( + ppid=0, + rss=1, + vms=2, + create_time=3, + nice=4, + num_threads=5, + status=6, + ttynr=7, +) + + +# ===================================================================== +# --- memory +# ===================================================================== + + +def virtual_memory(): + total, avail, free, _pinned, inuse = cext.virtual_mem() + percent = usage_percent((total - avail), total, round_=1) + return ntp.svmem(total, avail, percent, inuse, free) + + +def swap_memory(): + """Swap system memory as a (total, used, free, sin, sout) tuple.""" + total, free, sin, sout = cext.swap_mem() + used = total - free + percent = usage_percent(used, total, round_=1) + return ntp.sswap(total, used, free, percent, sin, sout) + + +# ===================================================================== +# --- CPU +# ===================================================================== + + +def cpu_times(): + """Return system-wide CPU times as a named tuple.""" + ret = cext.per_cpu_times() + return ntp.scputimes(*[sum(x) for x in zip(*ret)]) + + +def per_cpu_times(): + """Return system per-CPU times as a list of named tuples.""" + ret = cext.per_cpu_times() + return [ntp.scputimes(*x) for x in ret] + + +def cpu_count_logical(): + """Return the number of logical CPUs in the system.""" + try: + return os.sysconf("SC_NPROCESSORS_ONLN") + except ValueError: + # mimic os.cpu_count() behavior + return None + + +def cpu_count_cores(): + cmd = ["lsdev", "-Cc", "processor"] + p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + stdout, stderr = p.communicate() + stdout, stderr = (x.decode(sys.stdout.encoding) for x in (stdout, stderr)) + if p.returncode != 0: + msg = f"{cmd!r} command error\n{stderr}" + raise RuntimeError(msg) + processors = stdout.strip().splitlines() + return len(processors) or None + + +def cpu_stats(): + """Return various CPU stats as a named tuple.""" + ctx_switches, interrupts, soft_interrupts, syscalls = cext.cpu_stats() + return ntp.scpustats(ctx_switches, interrupts, soft_interrupts, syscalls) + + +# ===================================================================== +# --- disks +# ===================================================================== + + +disk_io_counters = cext.disk_io_counters +disk_usage = _psposix.disk_usage + + +def disk_partitions(all=False): + """Return system disk partitions.""" + # TODO - the filtering logic should be better checked so that + # it tries to reflect 'df' as much as possible + retlist = [] + partitions = cext.disk_partitions() + for partition in partitions: + device, mountpoint, fstype, opts = partition + if device == 'none': + device = '' + if not all: + # Differently from, say, Linux, we don't have a list of + # common fs types so the best we can do, AFAIK, is to + # filter by filesystem having a total size > 0. + if not disk_usage(mountpoint).total: + continue + ntuple = ntp.sdiskpart(device, mountpoint, fstype, opts) + retlist.append(ntuple) + return retlist + + +# ===================================================================== +# --- network +# ===================================================================== + + +net_if_addrs = cext.net_if_addrs + +if HAS_NET_IO_COUNTERS: + net_io_counters = cext.net_io_counters + + +def net_connections(kind, _pid=-1): + """Return socket connections. If pid == -1 return system-wide + connections (as opposed to connections opened by one process only). + """ + families, types = _common.conn_tmap[kind] + rawlist = cext.net_connections(_pid) + ret = [] + for item in rawlist: + fd, fam, type_, laddr, raddr, status, pid = item + if fam not in families: + continue + if type_ not in types: + continue + nt = conn_to_ntuple( + fd, + fam, + type_, + laddr, + raddr, + status, + TCP_STATUSES, + pid=pid if _pid == -1 else None, + ) + ret.append(nt) + return ret + + +def net_if_stats(): + """Get NIC stats (isup, duplex, speed, mtu).""" + duplex_map = {"Full": NIC_DUPLEX_FULL, "Half": NIC_DUPLEX_HALF} + names = {x[0] for x in net_if_addrs()} + ret = {} + for name in names: + mtu = cext.net_if_mtu(name) + flags = cext.net_if_flags(name) + + # try to get speed and duplex + # TODO: rewrite this in C (entstat forks, so use truss -f to follow. + # looks like it is using an undocumented ioctl?) + duplex = "" + speed = 0 + p = subprocess.Popen( + ["/usr/bin/entstat", "-d", name], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + stdout, stderr = p.communicate() + stdout, stderr = ( + x.decode(sys.stdout.encoding) for x in (stdout, stderr) + ) + if p.returncode == 0: + re_result = re.search( + r"Running: (\d+) Mbps.*?(\w+) Duplex", stdout + ) + if re_result is not None: + speed = int(re_result.group(1)) + duplex = re_result.group(2) + + output_flags = ','.join(flags) + isup = 'running' in flags + duplex = duplex_map.get(duplex, NIC_DUPLEX_UNKNOWN) + ret[name] = ntp.snicstats(isup, duplex, speed, mtu, output_flags) + return ret + + +# ===================================================================== +# --- other system functions +# ===================================================================== + + +def boot_time(): + """The system boot time expressed in seconds since the epoch.""" + return cext.boot_time() + + +def users(): + """Return currently connected users as a list of namedtuples.""" + retlist = [] + rawlist = cext.users() + localhost = (':0.0', ':0') + for item in rawlist: + user, tty, hostname, tstamp, user_process, pid = item + # note: the underlying C function includes entries about + # system boot, run level and others. We might want + # to use them in the future. + if not user_process: + continue + if hostname in localhost: + hostname = 'localhost' + nt = ntp.suser(user, tty, hostname, tstamp, pid) + retlist.append(nt) + return retlist + + +# ===================================================================== +# --- processes +# ===================================================================== + + +def pids(): + """Returns a list of PIDs currently running on the system.""" + return [int(x) for x in os.listdir(get_procfs_path()) if x.isdigit()] + + +def pid_exists(pid): + """Check for the existence of a unix pid.""" + return os.path.exists(os.path.join(get_procfs_path(), str(pid), "psinfo")) + + +def wrap_exceptions(fun): + """Call callable into a try/except clause and translate ENOENT, + EACCES and EPERM in NoSuchProcess or AccessDenied exceptions. + """ + + @functools.wraps(fun) + def wrapper(self, *args, **kwargs): + pid, ppid, name = self.pid, self._ppid, self._name + try: + return fun(self, *args, **kwargs) + except (FileNotFoundError, ProcessLookupError) as err: + # ENOENT (no such file or directory) gets raised on open(). + # ESRCH (no such process) can get raised on read() if + # process is gone in meantime. + if not pid_exists(pid): + raise NoSuchProcess(pid, name) from err + raise ZombieProcess(pid, name, ppid) from err + except PermissionError as err: + raise AccessDenied(pid, name) from err + + return wrapper + + +class Process: + """Wrapper class around underlying C implementation.""" + + __slots__ = ["_cache", "_name", "_ppid", "_procfs_path", "pid"] + + def __init__(self, pid): + self.pid = pid + self._name = None + self._ppid = None + self._procfs_path = get_procfs_path() + + def oneshot_enter(self): + self._proc_basic_info.cache_activate(self) + self._proc_cred.cache_activate(self) + + def oneshot_exit(self): + self._proc_basic_info.cache_deactivate(self) + self._proc_cred.cache_deactivate(self) + + @wrap_exceptions + @memoize_when_activated + def _proc_basic_info(self): + return cext.proc_basic_info(self.pid, self._procfs_path) + + @wrap_exceptions + @memoize_when_activated + def _proc_cred(self): + return cext.proc_cred(self.pid, self._procfs_path) + + @wrap_exceptions + def name(self): + if self.pid == 0: + return "swapper" + # note: max 16 characters + return cext.proc_name(self.pid, self._procfs_path).rstrip("\x00") + + @wrap_exceptions + def exe(self): + # there is no way to get executable path in AIX other than to guess, + # and guessing is more complex than what's in the wrapping class + cmdline = self.cmdline() + if not cmdline: + return '' + exe = cmdline[0] + if os.path.sep in exe: + # relative or absolute path + if not os.path.isabs(exe): + # if cwd has changed, we're out of luck - this may be wrong! + exe = os.path.abspath(os.path.join(self.cwd(), exe)) + if ( + os.path.isabs(exe) + and os.path.isfile(exe) + and os.access(exe, os.X_OK) + ): + return exe + # not found, move to search in PATH using basename only + exe = os.path.basename(exe) + # search for exe name PATH + for path in os.environ["PATH"].split(":"): + possible_exe = os.path.abspath(os.path.join(path, exe)) + if os.path.isfile(possible_exe) and os.access( + possible_exe, os.X_OK + ): + return possible_exe + return '' + + @wrap_exceptions + def cmdline(self): + return cext.proc_args(self.pid) + + @wrap_exceptions + def environ(self): + return cext.proc_environ(self.pid) + + @wrap_exceptions + def create_time(self): + return self._proc_basic_info()[proc_info_map['create_time']] + + @wrap_exceptions + def num_threads(self): + return self._proc_basic_info()[proc_info_map['num_threads']] + + if HAS_THREADS: + + @wrap_exceptions + def threads(self): + rawlist = cext.proc_threads(self.pid) + retlist = [] + for thread_id, utime, stime in rawlist: + ntuple = ntp.pthread(thread_id, utime, stime) + retlist.append(ntuple) + # The underlying C implementation retrieves all OS threads + # and filters them by PID. At this point we can't tell whether + # an empty list means there were no connections for process or + # process is no longer active so we force NSP in case the PID + # is no longer there. + if not retlist: + # will raise NSP if process is gone + os.stat(f"{self._procfs_path}/{self.pid}") + return retlist + + @wrap_exceptions + def net_connections(self, kind='inet'): + ret = net_connections(kind, _pid=self.pid) + # The underlying C implementation retrieves all OS connections + # and filters them by PID. At this point we can't tell whether + # an empty list means there were no connections for process or + # process is no longer active so we force NSP in case the PID + # is no longer there. + if not ret: + # will raise NSP if process is gone + os.stat(f"{self._procfs_path}/{self.pid}") + return ret + + @wrap_exceptions + def nice_get(self): + return cext.proc_priority_get(self.pid) + + @wrap_exceptions + def nice_set(self, value): + return cext.proc_priority_set(self.pid, value) + + @wrap_exceptions + def ppid(self): + self._ppid = self._proc_basic_info()[proc_info_map['ppid']] + return self._ppid + + @wrap_exceptions + def uids(self): + real, effective, saved, _, _, _ = self._proc_cred() + return ntp.puids(real, effective, saved) + + @wrap_exceptions + def gids(self): + _, _, _, real, effective, saved = self._proc_cred() + return ntp.puids(real, effective, saved) + + @wrap_exceptions + def cpu_times(self): + t = cext.proc_cpu_times(self.pid, self._procfs_path) + return ntp.pcputimes(*t) + + @wrap_exceptions + def terminal(self): + ttydev = self._proc_basic_info()[proc_info_map['ttynr']] + # convert from 64-bit dev_t to 32-bit dev_t and then map the device + ttydev = ((ttydev & 0x0000FFFF00000000) >> 16) | (ttydev & 0xFFFF) + # try to match rdev of /dev/pts/* files ttydev + for dev in glob.glob("/dev/**/*"): + if os.stat(dev).st_rdev == ttydev: + return dev + return None + + @wrap_exceptions + def cwd(self): + procfs_path = self._procfs_path + try: + result = os.readlink(f"{procfs_path}/{self.pid}/cwd") + return result.rstrip('/') + except FileNotFoundError: + os.stat(f"{procfs_path}/{self.pid}") # raise NSP or AD + return "" + + @wrap_exceptions + def memory_info(self): + ret = self._proc_basic_info() + rss = ret[proc_info_map['rss']] * 1024 + vms = ret[proc_info_map['vms']] * 1024 + return ntp.pmem(rss, vms) + + memory_full_info = memory_info + + @wrap_exceptions + def status(self): + code = self._proc_basic_info()[proc_info_map['status']] + # XXX is '?' legit? (we're not supposed to return it anyway) + return PROC_STATUSES.get(code, '?') + + def open_files(self): + # TODO rewrite without using procfiles (stat /proc/pid/fd/* and then + # find matching name of the inode) + p = subprocess.Popen( + ["/usr/bin/procfiles", "-n", str(self.pid)], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + stdout, stderr = p.communicate() + stdout, stderr = ( + x.decode(sys.stdout.encoding) for x in (stdout, stderr) + ) + if "no such process" in stderr.lower(): + raise NoSuchProcess(self.pid, self._name) + procfiles = re.findall(r"(\d+): S_IFREG.*name:(.*)\n", stdout) + retlist = [] + for fd, path in procfiles: + path = path.strip() + if path.startswith("//"): + path = path[1:] + if path.lower() == "cannot be retrieved": + continue + retlist.append(ntp.popenfile(path, int(fd))) + return retlist + + @wrap_exceptions + def num_fds(self): + if self.pid == 0: # no /proc/0/fd + return 0 + return len(os.listdir(f"{self._procfs_path}/{self.pid}/fd")) + + @wrap_exceptions + def num_ctx_switches(self): + return ntp.pctxsw(*cext.proc_num_ctx_switches(self.pid)) + + @wrap_exceptions + def wait(self, timeout=None): + return _psposix.wait_pid(self.pid, timeout) + + if HAS_PROC_IO_COUNTERS: + + @wrap_exceptions + def io_counters(self): + try: + rc, wc, rb, wb = cext.proc_io_counters(self.pid) + except OSError as err: + # if process is terminated, proc_io_counters returns OSError + # instead of NSP + if not pid_exists(self.pid): + raise NoSuchProcess(self.pid, self._name) from err + raise + return ntp.pio(rc, wc, rb, wb) diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/psutil/_psbsd.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/psutil/_psbsd.py new file mode 100644 index 0000000000000000000000000000000000000000..a75cd1af131973c0f78caa65a22c0a7e91645959 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/psutil/_psbsd.py @@ -0,0 +1,904 @@ +# Copyright (c) 2009, Giampaolo Rodola'. All rights reserved. +# Use of this source code is governed by a BSD-style license that can be +# found in the LICENSE file. + +"""FreeBSD, OpenBSD and NetBSD platforms implementation.""" + +import contextlib +import errno +import functools +import os +from collections import defaultdict +from collections import namedtuple +from xml.etree import ElementTree # noqa: ICN001 + +from . import _common +from . import _ntuples as ntp +from . import _psposix +from . import _psutil_bsd as cext +from ._common import FREEBSD +from ._common import NETBSD +from ._common import OPENBSD +from ._common import AccessDenied +from ._common import NoSuchProcess +from ._common import ZombieProcess +from ._common import conn_tmap +from ._common import conn_to_ntuple +from ._common import debug +from ._common import memoize +from ._common import memoize_when_activated +from ._common import usage_percent + +__extra__all__ = [] + + +# ===================================================================== +# --- globals +# ===================================================================== + + +if FREEBSD: + PROC_STATUSES = { + cext.SIDL: _common.STATUS_IDLE, + cext.SRUN: _common.STATUS_RUNNING, + cext.SSLEEP: _common.STATUS_SLEEPING, + cext.SSTOP: _common.STATUS_STOPPED, + cext.SZOMB: _common.STATUS_ZOMBIE, + cext.SWAIT: _common.STATUS_WAITING, + cext.SLOCK: _common.STATUS_LOCKED, + } +elif OPENBSD: + PROC_STATUSES = { + cext.SIDL: _common.STATUS_IDLE, + cext.SSLEEP: _common.STATUS_SLEEPING, + cext.SSTOP: _common.STATUS_STOPPED, + # According to /usr/include/sys/proc.h SZOMB is unused. + # test_zombie_process() shows that SDEAD is the right + # equivalent. Also it appears there's no equivalent of + # psutil.STATUS_DEAD. SDEAD really means STATUS_ZOMBIE. + # cext.SZOMB: _common.STATUS_ZOMBIE, + cext.SDEAD: _common.STATUS_ZOMBIE, + cext.SZOMB: _common.STATUS_ZOMBIE, + # From http://www.eecs.harvard.edu/~margo/cs161/videos/proc.h.txt + # OpenBSD has SRUN and SONPROC: SRUN indicates that a process + # is runnable but *not* yet running, i.e. is on a run queue. + # SONPROC indicates that the process is actually executing on + # a CPU, i.e. it is no longer on a run queue. + # As such we'll map SRUN to STATUS_WAKING and SONPROC to + # STATUS_RUNNING + cext.SRUN: _common.STATUS_WAKING, + cext.SONPROC: _common.STATUS_RUNNING, + } +elif NETBSD: + PROC_STATUSES = { + cext.SIDL: _common.STATUS_IDLE, + cext.SSLEEP: _common.STATUS_SLEEPING, + cext.SSTOP: _common.STATUS_STOPPED, + cext.SZOMB: _common.STATUS_ZOMBIE, + cext.SRUN: _common.STATUS_WAKING, + cext.SONPROC: _common.STATUS_RUNNING, + } + +TCP_STATUSES = { + cext.TCPS_ESTABLISHED: _common.CONN_ESTABLISHED, + cext.TCPS_SYN_SENT: _common.CONN_SYN_SENT, + cext.TCPS_SYN_RECEIVED: _common.CONN_SYN_RECV, + cext.TCPS_FIN_WAIT_1: _common.CONN_FIN_WAIT1, + cext.TCPS_FIN_WAIT_2: _common.CONN_FIN_WAIT2, + cext.TCPS_TIME_WAIT: _common.CONN_TIME_WAIT, + cext.TCPS_CLOSED: _common.CONN_CLOSE, + cext.TCPS_CLOSE_WAIT: _common.CONN_CLOSE_WAIT, + cext.TCPS_LAST_ACK: _common.CONN_LAST_ACK, + cext.TCPS_LISTEN: _common.CONN_LISTEN, + cext.TCPS_CLOSING: _common.CONN_CLOSING, + cext.PSUTIL_CONN_NONE: _common.CONN_NONE, +} + +PAGESIZE = cext.getpagesize() +AF_LINK = cext.AF_LINK + +HAS_PROC_NUM_THREADS = hasattr(cext, "proc_num_threads") + +kinfo_proc_map = dict( + ppid=0, + status=1, + real_uid=2, + effective_uid=3, + saved_uid=4, + real_gid=5, + effective_gid=6, + saved_gid=7, + ttynr=8, + create_time=9, + ctx_switches_vol=10, + ctx_switches_unvol=11, + read_io_count=12, + write_io_count=13, + user_time=14, + sys_time=15, + ch_user_time=16, + ch_sys_time=17, + rss=18, + vms=19, + memtext=20, + memdata=21, + memstack=22, + cpunum=23, + name=24, +) + + +# ===================================================================== +# --- memory +# ===================================================================== + + +def virtual_memory(): + mem = cext.virtual_mem() + if NETBSD: + total, free, active, inactive, wired, cached = mem + # On NetBSD buffers and shared mem is determined via /proc. + # The C ext set them to 0. + with open('/proc/meminfo', 'rb') as f: + for line in f: + if line.startswith(b'Buffers:'): + buffers = int(line.split()[1]) * 1024 + elif line.startswith(b'MemShared:'): + shared = int(line.split()[1]) * 1024 + # Before avail was calculated as (inactive + cached + free), + # same as zabbix, but it turned out it could exceed total (see + # #2233), so zabbix seems to be wrong. Htop calculates it + # differently, and the used value seem more realistic, so let's + # match htop. + # https://github.com/htop-dev/htop/blob/e7f447b/netbsd/NetBSDProcessList.c#L162 + # https://github.com/zabbix/zabbix/blob/af5e0f8/src/libs/zbxsysinfo/netbsd/memory.c#L135 + used = active + wired + avail = total - used + else: + total, free, active, inactive, wired, cached, buffers, shared = mem + # matches freebsd-memory CLI: + # * https://people.freebsd.org/~rse/dist/freebsd-memory + # * https://www.cyberciti.biz/files/scripts/freebsd-memory.pl.txt + # matches zabbix: + # * https://github.com/zabbix/zabbix/blob/af5e0f8/src/libs/zbxsysinfo/freebsd/memory.c#L143 + avail = inactive + cached + free + used = active + wired + cached + + percent = usage_percent((total - avail), total, round_=1) + return ntp.svmem( + total, + avail, + percent, + used, + free, + active, + inactive, + buffers, + cached, + shared, + wired, + ) + + +def swap_memory(): + """System swap memory as (total, used, free, sin, sout) namedtuple.""" + total, used, free, sin, sout = cext.swap_mem() + percent = usage_percent(used, total, round_=1) + return ntp.sswap(total, used, free, percent, sin, sout) + + +# malloc / heap functions (FreeBSD / NetBSD) +if hasattr(cext, "heap_info"): + heap_info = cext.heap_info + heap_trim = cext.heap_trim + + +# ===================================================================== +# --- CPU +# ===================================================================== + + +def cpu_times(): + """Return system per-CPU times as a namedtuple.""" + user, nice, system, idle, irq = cext.cpu_times() + return ntp.scputimes(user, nice, system, idle, irq) + + +def per_cpu_times(): + """Return system CPU times as a namedtuple.""" + ret = [] + for cpu_t in cext.per_cpu_times(): + user, nice, system, idle, irq = cpu_t + item = ntp.scputimes(user, nice, system, idle, irq) + ret.append(item) + return ret + + +def cpu_count_logical(): + """Return the number of logical CPUs in the system.""" + return cext.cpu_count_logical() + + +if OPENBSD or NETBSD: + + def cpu_count_cores(): + # OpenBSD and NetBSD do not implement this. + return 1 if cpu_count_logical() == 1 else None + +else: + + def cpu_count_cores(): + """Return the number of CPU cores in the system.""" + # From the C module we'll get an XML string similar to this: + # http://manpages.ubuntu.com/manpages/precise/man4/smp.4freebsd.html + # We may get None in case "sysctl kern.sched.topology_spec" + # is not supported on this BSD version, in which case we'll mimic + # os.cpu_count() and return None. + ret = None + s = cext.cpu_topology() + if s is not None: + # get rid of padding chars appended at the end of the string + index = s.rfind("") + if index != -1: + s = s[: index + 9] + root = ElementTree.fromstring(s) + try: + ret = len(root.findall('group/children/group/cpu')) or None + finally: + # needed otherwise it will memleak + root.clear() + if not ret: + # If logical CPUs == 1 it's obvious we' have only 1 core. + if cpu_count_logical() == 1: + return 1 + return ret + + +def cpu_stats(): + """Return various CPU stats as a named tuple.""" + if FREEBSD: + # Note: the C ext is returning some metrics we are not exposing: + # traps. + ctxsw, intrs, soft_intrs, syscalls, _traps = cext.cpu_stats() + elif NETBSD: + # XXX + # Note about intrs: the C extension returns 0. intrs + # can be determined via /proc/stat; it has the same value as + # soft_intrs thought so the kernel is faking it (?). + # + # Note about syscalls: the C extension always sets it to 0 (?). + # + # Note: the C ext is returning some metrics we are not exposing: + # traps, faults and forks. + ctxsw, intrs, soft_intrs, syscalls, _traps, _faults, _forks = ( + cext.cpu_stats() + ) + with open('/proc/stat', 'rb') as f: + for line in f: + if line.startswith(b'intr'): + intrs = int(line.split()[1]) + elif OPENBSD: + # Note: the C ext is returning some metrics we are not exposing: + # traps, faults and forks. + ctxsw, intrs, soft_intrs, syscalls, _traps, _faults, _forks = ( + cext.cpu_stats() + ) + return ntp.scpustats(ctxsw, intrs, soft_intrs, syscalls) + + +if FREEBSD: + + def cpu_freq(): + """Return frequency metrics for CPUs. As of Dec 2018 only + CPU 0 appears to be supported by FreeBSD and all other cores + match the frequency of CPU 0. + """ + ret = [] + num_cpus = cpu_count_logical() + for cpu in range(num_cpus): + try: + current, available_freq = cext.cpu_freq(cpu) + except NotImplementedError: + continue + if available_freq: + try: + min_freq = int(available_freq.split(" ")[-1].split("/")[0]) + except (IndexError, ValueError): + min_freq = None + try: + max_freq = int(available_freq.split(" ")[0].split("/")[0]) + except (IndexError, ValueError): + max_freq = None + ret.append(ntp.scpufreq(current, min_freq, max_freq)) + return ret + +elif OPENBSD: + + def cpu_freq(): + curr = float(cext.cpu_freq()) + return [ntp.scpufreq(curr, 0.0, 0.0)] + + +# ===================================================================== +# --- disks +# ===================================================================== + + +def disk_partitions(all=False): + """Return mounted disk partitions as a list of namedtuples. + 'all' argument is ignored, see: + https://github.com/giampaolo/psutil/issues/906. + """ + retlist = [] + partitions = cext.disk_partitions() + for partition in partitions: + device, mountpoint, fstype, opts = partition + ntuple = ntp.sdiskpart(device, mountpoint, fstype, opts) + retlist.append(ntuple) + return retlist + + +disk_usage = _psposix.disk_usage +disk_io_counters = cext.disk_io_counters + + +# ===================================================================== +# --- network +# ===================================================================== + + +net_io_counters = cext.net_io_counters +net_if_addrs = cext.net_if_addrs + + +def net_if_stats(): + """Get NIC stats (isup, duplex, speed, mtu).""" + names = net_io_counters().keys() + ret = {} + for name in names: + try: + mtu = cext.net_if_mtu(name) + flags = cext.net_if_flags(name) + duplex, speed = cext.net_if_duplex_speed(name) + except OSError as err: + # https://github.com/giampaolo/psutil/issues/1279 + if err.errno != errno.ENODEV: + raise + else: + if hasattr(_common, 'NicDuplex'): + duplex = _common.NicDuplex(duplex) + output_flags = ','.join(flags) + isup = 'running' in flags + ret[name] = ntp.snicstats(isup, duplex, speed, mtu, output_flags) + return ret + + +def net_connections(kind): + """System-wide network connections.""" + families, types = conn_tmap[kind] + ret = set() + if OPENBSD: + rawlist = cext.net_connections(-1, families, types) + elif NETBSD: + rawlist = cext.net_connections(-1, kind) + else: # FreeBSD + rawlist = cext.net_connections(families, types) + + for item in rawlist: + fd, fam, type, laddr, raddr, status, pid = item + nt = conn_to_ntuple( + fd, fam, type, laddr, raddr, status, TCP_STATUSES, pid + ) + ret.add(nt) + return list(ret) + + +# ===================================================================== +# --- sensors +# ===================================================================== + + +if FREEBSD: + + def sensors_battery(): + """Return battery info.""" + try: + percent, minsleft, power_plugged = cext.sensors_battery() + except NotImplementedError: + # See: https://github.com/giampaolo/psutil/issues/1074 + return None + power_plugged = power_plugged == 1 + if power_plugged: + secsleft = _common.POWER_TIME_UNLIMITED + elif minsleft == -1: + secsleft = _common.POWER_TIME_UNKNOWN + else: + secsleft = minsleft * 60 + return ntp.sbattery(percent, secsleft, power_plugged) + + def sensors_temperatures(): + """Return CPU cores temperatures if available, else an empty dict.""" + ret = defaultdict(list) + num_cpus = cpu_count_logical() + for cpu in range(num_cpus): + try: + current, high = cext.sensors_cpu_temperature(cpu) + if high <= 0: + high = None + name = f"Core {cpu}" + ret["coretemp"].append(ntp.shwtemp(name, current, high, high)) + except NotImplementedError: + pass + + return ret + + +# ===================================================================== +# --- other system functions +# ===================================================================== + + +def boot_time(): + """The system boot time expressed in seconds since the epoch.""" + return cext.boot_time() + + +if NETBSD: + + try: + INIT_BOOT_TIME = boot_time() + except Exception as err: # noqa: BLE001 + # Don't want to crash at import time. + debug(f"ignoring exception on import: {err!r}") + INIT_BOOT_TIME = 0 + + def adjust_proc_create_time(ctime): + """Account for system clock updates.""" + if INIT_BOOT_TIME == 0: + return ctime + + diff = INIT_BOOT_TIME - boot_time() + if diff == 0 or abs(diff) < 1: + return ctime + + debug("system clock was updated; adjusting process create_time()") + if diff < 0: + return ctime - diff + return ctime + diff + + +def users(): + """Return currently connected users as a list of namedtuples.""" + retlist = [] + rawlist = cext.users() + for item in rawlist: + user, tty, hostname, tstamp, pid = item + if tty == '~': + continue # reboot or shutdown + nt = ntp.suser(user, tty or None, hostname, tstamp, pid) + retlist.append(nt) + return retlist + + +# ===================================================================== +# --- processes +# ===================================================================== + + +@memoize +def _pid_0_exists(): + try: + Process(0).name() + except NoSuchProcess: + return False + except AccessDenied: + return True + else: + return True + + +def pids(): + """Returns a list of PIDs currently running on the system.""" + ret = cext.pids() + if OPENBSD and (0 not in ret) and _pid_0_exists(): + # On OpenBSD the kernel does not return PID 0 (neither does + # ps) but it's actually querable (Process(0) will succeed). + ret.insert(0, 0) + return ret + + +if NETBSD: + + def pid_exists(pid): + exists = _psposix.pid_exists(pid) + if not exists: + # We do this because _psposix.pid_exists() lies in case of + # zombie processes. + return pid in pids() + else: + return True + +elif OPENBSD: + + def pid_exists(pid): + exists = _psposix.pid_exists(pid) + if not exists: + return False + else: + # OpenBSD seems to be the only BSD platform where + # _psposix.pid_exists() returns True for thread IDs (tids), + # so we can't use it. + return pid in pids() + +else: # FreeBSD + pid_exists = _psposix.pid_exists + + +def wrap_exceptions(fun): + """Decorator which translates bare OSError exceptions into + NoSuchProcess and AccessDenied. + """ + + @functools.wraps(fun) + def wrapper(self, *args, **kwargs): + pid, ppid, name = self.pid, self._ppid, self._name + try: + return fun(self, *args, **kwargs) + except ProcessLookupError as err: + if cext.proc_is_zombie(pid): + raise ZombieProcess(pid, name, ppid) from err + raise NoSuchProcess(pid, name) from err + except PermissionError as err: + raise AccessDenied(pid, name) from err + except cext.ZombieProcessError as err: + raise ZombieProcess(pid, name, ppid) from err + except OSError as err: + if pid == 0 and 0 in pids(): + raise AccessDenied(pid, name) from err + raise err from None + + return wrapper + + +@contextlib.contextmanager +def wrap_exceptions_procfs(inst): + """Same as above, for routines relying on reading /proc fs.""" + pid, name, ppid = inst.pid, inst._name, inst._ppid + try: + yield + except (ProcessLookupError, FileNotFoundError) as err: + # ENOENT (no such file or directory) gets raised on open(). + # ESRCH (no such process) can get raised on read() if + # process is gone in meantime. + if cext.proc_is_zombie(inst.pid): + raise ZombieProcess(pid, name, ppid) from err + else: + raise NoSuchProcess(pid, name) from err + except PermissionError as err: + raise AccessDenied(pid, name) from err + + +class Process: + """Wrapper class around underlying C implementation.""" + + __slots__ = ["_cache", "_name", "_ppid", "pid"] + + def __init__(self, pid): + self.pid = pid + self._name = None + self._ppid = None + + def _assert_alive(self): + """Raise NSP if the process disappeared on us.""" + # For those C function who do not raise NSP, possibly returning + # incorrect or incomplete result. + cext.proc_name(self.pid) + + @wrap_exceptions + @memoize_when_activated + def oneshot(self): + """Retrieves multiple process info in one shot as a raw tuple.""" + ret = cext.proc_oneshot_info(self.pid) + assert len(ret) == len(kinfo_proc_map) + return ret + + def oneshot_enter(self): + self.oneshot.cache_activate(self) + + def oneshot_exit(self): + self.oneshot.cache_deactivate(self) + + @wrap_exceptions + def name(self): + name = self.oneshot()[kinfo_proc_map['name']] + return name if name is not None else cext.proc_name(self.pid) + + @wrap_exceptions + def exe(self): + if FREEBSD: + if self.pid == 0: + return '' # else NSP + return cext.proc_exe(self.pid) + elif NETBSD: + if self.pid == 0: + # /proc/0 dir exists but /proc/0/exe doesn't + return "" + with wrap_exceptions_procfs(self): + return os.readlink(f"/proc/{self.pid}/exe") + else: + # OpenBSD: exe cannot be determined; references: + # https://chromium.googlesource.com/chromium/src/base/+/ + # master/base_paths_posix.cc + # We try our best guess by using which against the first + # cmdline arg (may return None). + import shutil + + cmdline = self.cmdline() + if cmdline: + return shutil.which(cmdline[0]) or "" + else: + return "" + + @wrap_exceptions + def cmdline(self): + if OPENBSD and self.pid == 0: + return [] # ...else it crashes + elif NETBSD: + # XXX - most of the times the underlying sysctl() call on + # NetBSD and OpenBSD returns a truncated string. Also + # /proc/pid/cmdline behaves the same so it looks like this + # is a kernel bug. + try: + return cext.proc_cmdline(self.pid) + except OSError as err: + if err.errno == errno.EINVAL: + pid, name, ppid = self.pid, self._name, self._ppid + if cext.proc_is_zombie(self.pid): + raise ZombieProcess(pid, name, ppid) from err + if not pid_exists(self.pid): + raise NoSuchProcess(pid, name, ppid) from err + # XXX: this happens with unicode tests. It means the C + # routine is unable to decode invalid unicode chars. + debug(f"ignoring {err!r} and returning an empty list") + return [] + else: + raise + else: + return cext.proc_cmdline(self.pid) + + @wrap_exceptions + def environ(self): + return cext.proc_environ(self.pid) + + @wrap_exceptions + def terminal(self): + tty_nr = self.oneshot()[kinfo_proc_map['ttynr']] + tmap = _psposix.get_terminal_map() + try: + return tmap[tty_nr] + except KeyError: + return None + + @wrap_exceptions + def ppid(self): + self._ppid = self.oneshot()[kinfo_proc_map['ppid']] + return self._ppid + + @wrap_exceptions + def uids(self): + rawtuple = self.oneshot() + return ntp.puids( + rawtuple[kinfo_proc_map['real_uid']], + rawtuple[kinfo_proc_map['effective_uid']], + rawtuple[kinfo_proc_map['saved_uid']], + ) + + @wrap_exceptions + def gids(self): + rawtuple = self.oneshot() + return ntp.pgids( + rawtuple[kinfo_proc_map['real_gid']], + rawtuple[kinfo_proc_map['effective_gid']], + rawtuple[kinfo_proc_map['saved_gid']], + ) + + @wrap_exceptions + def cpu_times(self): + rawtuple = self.oneshot() + return ntp.pcputimes( + rawtuple[kinfo_proc_map['user_time']], + rawtuple[kinfo_proc_map['sys_time']], + rawtuple[kinfo_proc_map['ch_user_time']], + rawtuple[kinfo_proc_map['ch_sys_time']], + ) + + if FREEBSD: + + @wrap_exceptions + def cpu_num(self): + return self.oneshot()[kinfo_proc_map['cpunum']] + + @wrap_exceptions + def memory_info(self): + rawtuple = self.oneshot() + return ntp.pmem( + rawtuple[kinfo_proc_map['rss']], + rawtuple[kinfo_proc_map['vms']], + rawtuple[kinfo_proc_map['memtext']], + rawtuple[kinfo_proc_map['memdata']], + rawtuple[kinfo_proc_map['memstack']], + ) + + memory_full_info = memory_info + + @wrap_exceptions + def create_time(self, monotonic=False): + ctime = self.oneshot()[kinfo_proc_map['create_time']] + if NETBSD and not monotonic: + # NetBSD: ctime subject to system clock updates. + ctime = adjust_proc_create_time(ctime) + return ctime + + @wrap_exceptions + def num_threads(self): + if HAS_PROC_NUM_THREADS: + # FreeBSD / NetBSD + return cext.proc_num_threads(self.pid) + else: + return len(self.threads()) + + @wrap_exceptions + def num_ctx_switches(self): + rawtuple = self.oneshot() + return ntp.pctxsw( + rawtuple[kinfo_proc_map['ctx_switches_vol']], + rawtuple[kinfo_proc_map['ctx_switches_unvol']], + ) + + @wrap_exceptions + def threads(self): + # Note: on OpenSBD this (/dev/mem) requires root access. + rawlist = cext.proc_threads(self.pid) + retlist = [] + for thread_id, utime, stime in rawlist: + ntuple = ntp.pthread(thread_id, utime, stime) + retlist.append(ntuple) + if OPENBSD: + self._assert_alive() + return retlist + + @wrap_exceptions + def net_connections(self, kind='inet'): + families, types = conn_tmap[kind] + ret = [] + + if NETBSD: + rawlist = cext.net_connections(self.pid, kind) + elif OPENBSD: + rawlist = cext.net_connections(self.pid, families, types) + else: + rawlist = cext.proc_net_connections(self.pid, families, types) + + for item in rawlist: + fd, fam, type, laddr, raddr, status = item[:6] + if FREEBSD: + if (fam not in families) or (type not in types): + continue + nt = conn_to_ntuple( + fd, fam, type, laddr, raddr, status, TCP_STATUSES + ) + ret.append(nt) + + self._assert_alive() + return ret + + @wrap_exceptions + def wait(self, timeout=None): + return _psposix.wait_pid(self.pid, timeout) + + @wrap_exceptions + def nice_get(self): + return cext.proc_priority_get(self.pid) + + @wrap_exceptions + def nice_set(self, value): + return cext.proc_priority_set(self.pid, value) + + @wrap_exceptions + def status(self): + code = self.oneshot()[kinfo_proc_map['status']] + # XXX is '?' legit? (we're not supposed to return it anyway) + return PROC_STATUSES.get(code, '?') + + @wrap_exceptions + def io_counters(self): + rawtuple = self.oneshot() + return ntp.pio( + rawtuple[kinfo_proc_map['read_io_count']], + rawtuple[kinfo_proc_map['write_io_count']], + -1, + -1, + ) + + @wrap_exceptions + def cwd(self): + """Return process current working directory.""" + # sometimes we get an empty string, in which case we turn + # it into None + if OPENBSD and self.pid == 0: + return "" # ...else it would raise EINVAL + return cext.proc_cwd(self.pid) + + nt_mmap_grouped = namedtuple( + 'mmap', 'path rss, private, ref_count, shadow_count' + ) + nt_mmap_ext = namedtuple( + 'mmap', 'addr, perms path rss, private, ref_count, shadow_count' + ) + + @wrap_exceptions + def open_files(self): + """Return files opened by process as a list of namedtuples.""" + rawlist = cext.proc_open_files(self.pid) + return [ntp.popenfile(path, fd) for path, fd in rawlist] + + @wrap_exceptions + def num_fds(self): + """Return the number of file descriptors opened by this process.""" + ret = cext.proc_num_fds(self.pid) + if NETBSD: + self._assert_alive() + return ret + + # --- FreeBSD only APIs + + if FREEBSD: + + @wrap_exceptions + def cpu_affinity_get(self): + return cext.proc_cpu_affinity_get(self.pid) + + @wrap_exceptions + def cpu_affinity_set(self, cpus): + # Pre-emptively check if CPUs are valid because the C + # function has a weird behavior in case of invalid CPUs, + # see: https://github.com/giampaolo/psutil/issues/586 + allcpus = set(range(len(per_cpu_times()))) + for cpu in cpus: + if cpu not in allcpus: + msg = f"invalid CPU {cpu!r} (choose between {allcpus})" + raise ValueError(msg) + try: + cext.proc_cpu_affinity_set(self.pid, cpus) + except OSError as err: + # 'man cpuset_setaffinity' about EDEADLK: + # <> + if err.errno in {errno.EINVAL, errno.EDEADLK}: + for cpu in cpus: + if cpu not in allcpus: + msg = ( + f"invalid CPU {cpu!r} (choose between" + f" {allcpus})" + ) + raise ValueError(msg) from err + raise + + @wrap_exceptions + def memory_maps(self): + return cext.proc_memory_maps(self.pid) + + @wrap_exceptions + def rlimit(self, resource, limits=None): + if limits is None: + return cext.proc_getrlimit(self.pid, resource) + else: + if len(limits) != 2: + msg = ( + "second argument must be a (soft, hard) tuple, got" + f" {limits!r}" + ) + raise ValueError(msg) + soft, hard = limits + return cext.proc_setrlimit(self.pid, resource, soft, hard) diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/psutil/_pslinux.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/psutil/_pslinux.py new file mode 100644 index 0000000000000000000000000000000000000000..dc305b59559e9b97f96eacf2a2002d2b02e3fe41 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/psutil/_pslinux.py @@ -0,0 +1,2269 @@ +# Copyright (c) 2009, Giampaolo Rodola'. All rights reserved. +# Use of this source code is governed by a BSD-style license that can be +# found in the LICENSE file. + +"""Linux platform implementation.""" + +import base64 +import collections +import enum +import errno +import functools +import glob +import os +import re +import resource +import socket +import struct +import sys +import warnings +from collections import defaultdict +from collections import namedtuple + +from . import _common +from . import _ntuples as ntp +from . import _psposix +from . import _psutil_linux as cext +from ._common import ENCODING +from ._common import NIC_DUPLEX_FULL +from ._common import NIC_DUPLEX_HALF +from ._common import NIC_DUPLEX_UNKNOWN +from ._common import AccessDenied +from ._common import NoSuchProcess +from ._common import ZombieProcess +from ._common import bcat +from ._common import cat +from ._common import debug +from ._common import decode +from ._common import get_procfs_path +from ._common import isfile_strict +from ._common import memoize +from ._common import memoize_when_activated +from ._common import open_binary +from ._common import open_text +from ._common import parse_environ_block +from ._common import path_exists_strict +from ._common import supports_ipv6 +from ._common import usage_percent + +# fmt: off +__extra__all__ = [ + 'PROCFS_PATH', + # io prio constants + "IOPRIO_CLASS_NONE", "IOPRIO_CLASS_RT", "IOPRIO_CLASS_BE", + "IOPRIO_CLASS_IDLE", + # connection status constants + "CONN_ESTABLISHED", "CONN_SYN_SENT", "CONN_SYN_RECV", "CONN_FIN_WAIT1", + "CONN_FIN_WAIT2", "CONN_TIME_WAIT", "CONN_CLOSE", "CONN_CLOSE_WAIT", + "CONN_LAST_ACK", "CONN_LISTEN", "CONN_CLOSING", +] +# fmt: on + + +# ===================================================================== +# --- globals +# ===================================================================== + + +POWER_SUPPLY_PATH = "/sys/class/power_supply" +HAS_PROC_SMAPS = os.path.exists(f"/proc/{os.getpid()}/smaps") +HAS_PROC_SMAPS_ROLLUP = os.path.exists(f"/proc/{os.getpid()}/smaps_rollup") +HAS_PROC_IO_PRIORITY = hasattr(cext, "proc_ioprio_get") +HAS_CPU_AFFINITY = hasattr(cext, "proc_cpu_affinity_get") + +# Number of clock ticks per second +CLOCK_TICKS = os.sysconf("SC_CLK_TCK") +PAGESIZE = cext.getpagesize() +LITTLE_ENDIAN = sys.byteorder == 'little' +UNSET = object() + +# "man iostat" states that sectors are equivalent with blocks and have +# a size of 512 bytes. Despite this value can be queried at runtime +# via /sys/block/{DISK}/queue/hw_sector_size and results may vary +# between 1k, 2k, or 4k... 512 appears to be a magic constant used +# throughout Linux source code: +# * https://stackoverflow.com/a/38136179/376587 +# * https://lists.gt.net/linux/kernel/2241060 +# * https://github.com/giampaolo/psutil/issues/1305 +# * https://github.com/torvalds/linux/blob/ +# 4f671fe2f9523a1ea206f63fe60a7c7b3a56d5c7/include/linux/bio.h#L99 +# * https://lkml.org/lkml/2015/8/17/234 +DISK_SECTOR_SIZE = 512 + +AddressFamily = enum.IntEnum( + 'AddressFamily', {'AF_LINK': int(socket.AF_PACKET)} +) +AF_LINK = AddressFamily.AF_LINK + + +# ioprio_* constants http://linux.die.net/man/2/ioprio_get +class IOPriority(enum.IntEnum): + IOPRIO_CLASS_NONE = 0 + IOPRIO_CLASS_RT = 1 + IOPRIO_CLASS_BE = 2 + IOPRIO_CLASS_IDLE = 3 + + +globals().update(IOPriority.__members__) + +# See: +# https://github.com/torvalds/linux/blame/master/fs/proc/array.c +# ...and (TASK_* constants): +# https://github.com/torvalds/linux/blob/master/include/linux/sched.h +PROC_STATUSES = { + "R": _common.STATUS_RUNNING, + "S": _common.STATUS_SLEEPING, + "D": _common.STATUS_DISK_SLEEP, + "T": _common.STATUS_STOPPED, + "t": _common.STATUS_TRACING_STOP, + "Z": _common.STATUS_ZOMBIE, + "X": _common.STATUS_DEAD, + "x": _common.STATUS_DEAD, + "K": _common.STATUS_WAKE_KILL, + "W": _common.STATUS_WAKING, + "I": _common.STATUS_IDLE, + "P": _common.STATUS_PARKED, +} + +# https://github.com/torvalds/linux/blob/master/include/net/tcp_states.h +TCP_STATUSES = { + "01": _common.CONN_ESTABLISHED, + "02": _common.CONN_SYN_SENT, + "03": _common.CONN_SYN_RECV, + "04": _common.CONN_FIN_WAIT1, + "05": _common.CONN_FIN_WAIT2, + "06": _common.CONN_TIME_WAIT, + "07": _common.CONN_CLOSE, + "08": _common.CONN_CLOSE_WAIT, + "09": _common.CONN_LAST_ACK, + "0A": _common.CONN_LISTEN, + "0B": _common.CONN_CLOSING, +} + + +# ===================================================================== +# --- utils +# ===================================================================== + + +def readlink(path): + """Wrapper around os.readlink().""" + assert isinstance(path, str), path + path = os.readlink(path) + # readlink() might return paths containing null bytes ('\x00') + # resulting in "TypeError: must be encoded string without NULL + # bytes, not str" errors when the string is passed to other + # fs-related functions (os.*, open(), ...). + # Apparently everything after '\x00' is garbage (we can have + # ' (deleted)', 'new' and possibly others), see: + # https://github.com/giampaolo/psutil/issues/717 + path = path.split('\x00')[0] + # Certain paths have ' (deleted)' appended. Usually this is + # bogus as the file actually exists. Even if it doesn't we + # don't care. + if path.endswith(' (deleted)') and not path_exists_strict(path): + path = path[:-10] + return path + + +def file_flags_to_mode(flags): + """Convert file's open() flags into a readable string. + Used by Process.open_files(). + """ + modes_map = {os.O_RDONLY: 'r', os.O_WRONLY: 'w', os.O_RDWR: 'w+'} + mode = modes_map[flags & (os.O_RDONLY | os.O_WRONLY | os.O_RDWR)] + if flags & os.O_APPEND: + mode = mode.replace('w', 'a', 1) + mode = mode.replace('w+', 'r+') + # possible values: r, w, a, r+, a+ + return mode + + +def is_storage_device(name): + """Return True if the given name refers to a root device (e.g. + "sda", "nvme0n1") as opposed to a logical partition (e.g. "sda1", + "nvme0n1p1"). If name is a virtual device (e.g. "loop1", "ram") + return True. + """ + # Re-adapted from iostat source code, see: + # https://github.com/sysstat/sysstat/blob/ + # 97912938cd476645b267280069e83b1c8dc0e1c7/common.c#L208 + # Some devices may have a slash in their name (e.g. cciss/c0d0...). + name = name.replace('/', '!') + including_virtual = True + if including_virtual: + path = f"/sys/block/{name}" + else: + path = f"/sys/block/{name}/device" + return os.access(path, os.F_OK) + + +@memoize +def _scputimes_ntuple(procfs_path): + """Return a namedtuple of variable fields depending on the CPU times + available on this Linux kernel version which may be: + (user, nice, system, idle, iowait, irq, softirq, [steal, [guest, + [guest_nice]]]) + Used by cpu_times() function. + """ + with open_binary(f"{procfs_path}/stat") as f: + values = f.readline().split()[1:] + fields = ['user', 'nice', 'system', 'idle', 'iowait', 'irq', 'softirq'] + vlen = len(values) + if vlen >= 8: + # Linux >= 2.6.11 + fields.append('steal') + if vlen >= 9: + # Linux >= 2.6.24 + fields.append('guest') + if vlen >= 10: + # Linux >= 3.2.0 + fields.append('guest_nice') + return namedtuple('scputimes', fields) + + +# Set it into _ntuples.py namespace. +try: + ntp.scputimes = _scputimes_ntuple("/proc") +except Exception as err: # noqa: BLE001 + # Don't want to crash at import time. + debug(f"ignoring exception on import: {err!r}") + ntp.scputimes = namedtuple('scputimes', 'user system idle')(0.0, 0.0, 0.0) + +# XXX: must be available also at this module level in order to be +# serialized (tests/test_misc.py::TestMisc::test_serialization). +scputimes = ntp.scputimes + + +# ===================================================================== +# --- system memory +# ===================================================================== + + +def calculate_avail_vmem(mems): + """Fallback for kernels < 3.14 where /proc/meminfo does not provide + "MemAvailable", see: + https://blog.famzah.net/2014/09/24/. + + This code reimplements the algorithm outlined here: + https://git.kernel.org/cgit/linux/kernel/git/torvalds/linux.git/ + commit/?id=34e431b0ae398fc54ea69ff85ec700722c9da773 + + We use this function also when "MemAvailable" returns 0 (possibly a + kernel bug, see: https://github.com/giampaolo/psutil/issues/1915). + In that case this routine matches "free" CLI tool result ("available" + column). + + XXX: on recent kernels this calculation may differ by ~1.5% compared + to "MemAvailable:", as it's calculated slightly differently. + It is still way more realistic than doing (free + cached) though. + See: + * https://gitlab.com/procps-ng/procps/issues/42 + * https://github.com/famzah/linux-memavailable-procfs/issues/2 + """ + # Note about "fallback" value. According to: + # https://git.kernel.org/cgit/linux/kernel/git/torvalds/linux.git/ + # commit/?id=34e431b0ae398fc54ea69ff85ec700722c9da773 + # ...long ago "available" memory was calculated as (free + cached), + # We use fallback when one of these is missing from /proc/meminfo: + # "Active(file)": introduced in 2.6.28 / Dec 2008 + # "Inactive(file)": introduced in 2.6.28 / Dec 2008 + # "SReclaimable": introduced in 2.6.19 / Nov 2006 + # /proc/zoneinfo: introduced in 2.6.13 / Aug 2005 + free = mems[b'MemFree:'] + fallback = free + mems.get(b"Cached:", 0) + try: + lru_active_file = mems[b'Active(file):'] + lru_inactive_file = mems[b'Inactive(file):'] + slab_reclaimable = mems[b'SReclaimable:'] + except KeyError as err: + debug( + f"{err.args[0]} is missing from /proc/meminfo; using an" + " approximation for calculating available memory" + ) + return fallback + try: + f = open_binary(f"{get_procfs_path()}/zoneinfo") + except OSError: + return fallback # kernel 2.6.13 + + watermark_low = 0 + with f: + for line in f: + line = line.strip() + if line.startswith(b'low'): + watermark_low += int(line.split()[1]) + watermark_low *= PAGESIZE + + avail = free - watermark_low + pagecache = lru_active_file + lru_inactive_file + pagecache -= min(pagecache / 2, watermark_low) + avail += pagecache + avail += slab_reclaimable - min(slab_reclaimable / 2.0, watermark_low) + return int(avail) + + +def virtual_memory(): + """Report virtual memory stats. + This implementation mimics procps-ng-3.3.12, aka "free" CLI tool: + https://gitlab.com/procps-ng/procps/blob/ + 24fd2605c51fccc375ab0287cec33aa767f06718/proc/sysinfo.c#L778-791 + The returned values are supposed to match both "free" and "vmstat -s" + CLI tools. + """ + missing_fields = [] + mems = {} + with open_binary(f"{get_procfs_path()}/meminfo") as f: + for line in f: + fields = line.split() + mems[fields[0]] = int(fields[1]) * 1024 + + # /proc doc states that the available fields in /proc/meminfo vary + # by architecture and compile options, but these 3 values are also + # returned by sysinfo(2); as such we assume they are always there. + total = mems[b'MemTotal:'] + free = mems[b'MemFree:'] + try: + buffers = mems[b'Buffers:'] + except KeyError: + # https://github.com/giampaolo/psutil/issues/1010 + buffers = 0 + missing_fields.append('buffers') + try: + cached = mems[b"Cached:"] + except KeyError: + cached = 0 + missing_fields.append('cached') + else: + # "free" cmdline utility sums reclaimable to cached. + # Older versions of procps used to add slab memory instead. + # This got changed in: + # https://gitlab.com/procps-ng/procps/commit/ + # 05d751c4f076a2f0118b914c5e51cfbb4762ad8e + cached += mems.get(b"SReclaimable:", 0) # since kernel 2.6.19 + + try: + shared = mems[b'Shmem:'] # since kernel 2.6.32 + except KeyError: + try: + shared = mems[b'MemShared:'] # kernels 2.4 + except KeyError: + shared = 0 + missing_fields.append('shared') + + try: + active = mems[b"Active:"] + except KeyError: + active = 0 + missing_fields.append('active') + + try: + inactive = mems[b"Inactive:"] + except KeyError: + try: + inactive = ( + mems[b"Inact_dirty:"] + + mems[b"Inact_clean:"] + + mems[b"Inact_laundry:"] + ) + except KeyError: + inactive = 0 + missing_fields.append('inactive') + + try: + slab = mems[b"Slab:"] + except KeyError: + slab = 0 + + # - starting from 4.4.0 we match free's "available" column. + # Before 4.4.0 we calculated it as (free + buffers + cached) + # which matched htop. + # - free and htop available memory differs as per: + # http://askubuntu.com/a/369589 + # http://unix.stackexchange.com/a/65852/168884 + # - MemAvailable has been introduced in kernel 3.14 + try: + avail = mems[b'MemAvailable:'] + except KeyError: + avail = calculate_avail_vmem(mems) + else: + if avail == 0: + # Yes, it can happen (probably a kernel bug): + # https://github.com/giampaolo/psutil/issues/1915 + # In this case "free" CLI tool makes an estimate. We do the same, + # and it matches "free" CLI tool. + avail = calculate_avail_vmem(mems) + + if avail < 0: + avail = 0 + missing_fields.append('available') + elif avail > total: + # If avail is greater than total or our calculation overflows, + # that's symptomatic of running within a LCX container where such + # values will be dramatically distorted over those of the host. + # https://gitlab.com/procps-ng/procps/blob/ + # 24fd2605c51fccc375ab0287cec33aa767f06718/proc/sysinfo.c#L764 + avail = free + + used = total - avail + + percent = usage_percent((total - avail), total, round_=1) + + # Warn about missing metrics which are set to 0. + if missing_fields: + msg = "{} memory stats couldn't be determined and {} set to 0".format( + ", ".join(missing_fields), + "was" if len(missing_fields) == 1 else "were", + ) + warnings.warn(msg, RuntimeWarning, stacklevel=2) + + return ntp.svmem( + total, + avail, + percent, + used, + free, + active, + inactive, + buffers, + cached, + shared, + slab, + ) + + +def swap_memory(): + """Return swap memory metrics.""" + mems = {} + with open_binary(f"{get_procfs_path()}/meminfo") as f: + for line in f: + fields = line.split() + mems[fields[0]] = int(fields[1]) * 1024 + # We prefer /proc/meminfo over sysinfo() syscall so that + # psutil.PROCFS_PATH can be used in order to allow retrieval + # for linux containers, see: + # https://github.com/giampaolo/psutil/issues/1015 + try: + total = mems[b'SwapTotal:'] + free = mems[b'SwapFree:'] + except KeyError: + _, _, _, _, total, free, unit_multiplier = cext.linux_sysinfo() + total *= unit_multiplier + free *= unit_multiplier + + used = total - free + percent = usage_percent(used, total, round_=1) + # get pgin/pgouts + try: + f = open_binary(f"{get_procfs_path()}/vmstat") + except OSError as err: + # see https://github.com/giampaolo/psutil/issues/722 + msg = ( + "'sin' and 'sout' swap memory stats couldn't " + f"be determined and were set to 0 ({err})" + ) + warnings.warn(msg, RuntimeWarning, stacklevel=2) + sin = sout = 0 + else: + with f: + sin = sout = None + for line in f: + # values are expressed in 4 kilo bytes, we want + # bytes instead + if line.startswith(b'pswpin'): + sin = int(line.split(b' ')[1]) * 4 * 1024 + elif line.startswith(b'pswpout'): + sout = int(line.split(b' ')[1]) * 4 * 1024 + if sin is not None and sout is not None: + break + else: + # we might get here when dealing with exotic Linux + # flavors, see: + # https://github.com/giampaolo/psutil/issues/313 + msg = "'sin' and 'sout' swap memory stats couldn't " + msg += "be determined and were set to 0" + warnings.warn(msg, RuntimeWarning, stacklevel=2) + sin = sout = 0 + return ntp.sswap(total, used, free, percent, sin, sout) + + +# malloc / heap functions; require glibc +if hasattr(cext, "heap_info"): + heap_info = cext.heap_info + heap_trim = cext.heap_trim + + +# ===================================================================== +# --- CPU +# ===================================================================== + + +def cpu_times(): + """Return a named tuple representing the following system-wide + CPU times: + (user, nice, system, idle, iowait, irq, softirq [steal, [guest, + [guest_nice]]]) + Last 3 fields may not be available on all Linux kernel versions. + """ + procfs_path = get_procfs_path() + with open_binary(f"{procfs_path}/stat") as f: + values = f.readline().split() + fields = values[1 : len(ntp.scputimes._fields) + 1] + fields = [float(x) / CLOCK_TICKS for x in fields] + return ntp.scputimes(*fields) + + +def per_cpu_times(): + """Return a list of namedtuple representing the CPU times + for every CPU available on the system. + """ + procfs_path = get_procfs_path() + cpus = [] + with open_binary(f"{procfs_path}/stat") as f: + # get rid of the first line which refers to system wide CPU stats + f.readline() + for line in f: + if line.startswith(b'cpu'): + values = line.split() + fields = values[1 : len(ntp.scputimes._fields) + 1] + fields = [float(x) / CLOCK_TICKS for x in fields] + entry = ntp.scputimes(*fields) + cpus.append(entry) + return cpus + + +def cpu_count_logical(): + """Return the number of logical CPUs in the system.""" + try: + return os.sysconf("SC_NPROCESSORS_ONLN") + except ValueError: + # as a second fallback we try to parse /proc/cpuinfo + num = 0 + with open_binary(f"{get_procfs_path()}/cpuinfo") as f: + for line in f: + if line.lower().startswith(b'processor'): + num += 1 + + # unknown format (e.g. amrel/sparc architectures), see: + # https://github.com/giampaolo/psutil/issues/200 + # try to parse /proc/stat as a last resort + if num == 0: + search = re.compile(r'cpu\d') + with open_text(f"{get_procfs_path()}/stat") as f: + for line in f: + line = line.split(' ')[0] + if search.match(line): + num += 1 + + if num == 0: + # mimic os.cpu_count() + return None + return num + + +def cpu_count_cores(): + """Return the number of CPU cores in the system.""" + # Method #1 + ls = set() + # These 2 files are the same but */core_cpus_list is newer while + # */thread_siblings_list is deprecated and may disappear in the future. + # https://www.kernel.org/doc/Documentation/admin-guide/cputopology.rst + # https://github.com/giampaolo/psutil/pull/1727#issuecomment-707624964 + # https://lkml.org/lkml/2019/2/26/41 + p1 = "/sys/devices/system/cpu/cpu[0-9]*/topology/core_cpus_list" + p2 = "/sys/devices/system/cpu/cpu[0-9]*/topology/thread_siblings_list" + for path in glob.glob(p1) or glob.glob(p2): + with open_binary(path) as f: + ls.add(f.read().strip()) + result = len(ls) + if result != 0: + return result + + # Method #2 + mapping = {} + current_info = {} + with open_binary(f"{get_procfs_path()}/cpuinfo") as f: + for line in f: + line = line.strip().lower() + if not line: + # new section + try: + mapping[current_info[b'physical id']] = current_info[ + b'cpu cores' + ] + except KeyError: + pass + current_info = {} + elif line.startswith((b'physical id', b'cpu cores')): + # ongoing section + key, value = line.split(b'\t:', 1) + current_info[key] = int(value) + + result = sum(mapping.values()) + return result or None # mimic os.cpu_count() + + +def cpu_stats(): + """Return various CPU stats as a named tuple.""" + with open_binary(f"{get_procfs_path()}/stat") as f: + ctx_switches = None + interrupts = None + soft_interrupts = None + for line in f: + if line.startswith(b'ctxt'): + ctx_switches = int(line.split()[1]) + elif line.startswith(b'intr'): + interrupts = int(line.split()[1]) + elif line.startswith(b'softirq'): + soft_interrupts = int(line.split()[1]) + if ( + ctx_switches is not None + and soft_interrupts is not None + and interrupts is not None + ): + break + syscalls = 0 + return ntp.scpustats(ctx_switches, interrupts, soft_interrupts, syscalls) + + +def _cpu_get_cpuinfo_freq(): + """Return current CPU frequency from cpuinfo if available.""" + with open_binary(f"{get_procfs_path()}/cpuinfo") as f: + return [ + float(line.split(b':', 1)[1]) + for line in f + if line.lower().startswith(b'cpu mhz') + ] + + +if os.path.exists("/sys/devices/system/cpu/cpufreq/policy0") or os.path.exists( + "/sys/devices/system/cpu/cpu0/cpufreq" +): + + def cpu_freq(): + """Return frequency metrics for all CPUs. + Contrarily to other OSes, Linux updates these values in + real-time. + """ + cpuinfo_freqs = _cpu_get_cpuinfo_freq() + paths = glob.glob( + "/sys/devices/system/cpu/cpufreq/policy[0-9]*" + ) or glob.glob("/sys/devices/system/cpu/cpu[0-9]*/cpufreq") + paths.sort(key=lambda x: int(re.search(r"[0-9]+", x).group())) + ret = [] + pjoin = os.path.join + for i, path in enumerate(paths): + if len(paths) == len(cpuinfo_freqs): + # take cached value from cpuinfo if available, see: + # https://github.com/giampaolo/psutil/issues/1851 + curr = cpuinfo_freqs[i] * 1000 + else: + curr = bcat(pjoin(path, "scaling_cur_freq"), fallback=None) + if curr is None: + # Likely an old RedHat, see: + # https://github.com/giampaolo/psutil/issues/1071 + curr = bcat(pjoin(path, "cpuinfo_cur_freq"), fallback=None) + if curr is None: + online_path = f"/sys/devices/system/cpu/cpu{i}/online" + # if cpu core is offline, set to all zeroes + if cat(online_path, fallback=None) == "0\n": + ret.append(ntp.scpufreq(0.0, 0.0, 0.0)) + continue + msg = "can't find current frequency file" + raise NotImplementedError(msg) + curr = int(curr) / 1000 + max_ = int(bcat(pjoin(path, "scaling_max_freq"))) / 1000 + min_ = int(bcat(pjoin(path, "scaling_min_freq"))) / 1000 + ret.append(ntp.scpufreq(curr, min_, max_)) + return ret + +else: + + def cpu_freq(): + """Alternate implementation using /proc/cpuinfo. + min and max frequencies are not available and are set to None. + """ + return [ntp.scpufreq(x, 0.0, 0.0) for x in _cpu_get_cpuinfo_freq()] + + +# ===================================================================== +# --- network +# ===================================================================== + + +net_if_addrs = cext.net_if_addrs + + +class _Ipv6UnsupportedError(Exception): + pass + + +class NetConnections: + """A wrapper on top of /proc/net/* files, retrieving per-process + and system-wide open connections (TCP, UDP, UNIX) similarly to + "netstat -an". + + Note: in case of UNIX sockets we're only able to determine the + local endpoint/path, not the one it's connected to. + According to [1] it would be possible but not easily. + + [1] http://serverfault.com/a/417946 + """ + + def __init__(self): + # The string represents the basename of the corresponding + # /proc/net/{proto_name} file. + tcp4 = ("tcp", socket.AF_INET, socket.SOCK_STREAM) + tcp6 = ("tcp6", socket.AF_INET6, socket.SOCK_STREAM) + udp4 = ("udp", socket.AF_INET, socket.SOCK_DGRAM) + udp6 = ("udp6", socket.AF_INET6, socket.SOCK_DGRAM) + unix = ("unix", socket.AF_UNIX, None) + self.tmap = { + "all": (tcp4, tcp6, udp4, udp6, unix), + "tcp": (tcp4, tcp6), + "tcp4": (tcp4,), + "tcp6": (tcp6,), + "udp": (udp4, udp6), + "udp4": (udp4,), + "udp6": (udp6,), + "unix": (unix,), + "inet": (tcp4, tcp6, udp4, udp6), + "inet4": (tcp4, udp4), + "inet6": (tcp6, udp6), + } + self._procfs_path = None + + def get_proc_inodes(self, pid): + inodes = defaultdict(list) + for fd in os.listdir(f"{self._procfs_path}/{pid}/fd"): + try: + inode = readlink(f"{self._procfs_path}/{pid}/fd/{fd}") + except (FileNotFoundError, ProcessLookupError): + # ENOENT == file which is gone in the meantime; + # os.stat(f"/proc/{self.pid}") will be done later + # to force NSP (if it's the case) + continue + except OSError as err: + if err.errno == errno.EINVAL: + # not a link + continue + if err.errno == errno.ENAMETOOLONG: + # file name too long + debug(err) + continue + raise + else: + if inode.startswith('socket:['): + # the process is using a socket + inode = inode[8:][:-1] + inodes[inode].append((pid, int(fd))) + return inodes + + def get_all_inodes(self): + inodes = {} + for pid in pids(): + try: + inodes.update(self.get_proc_inodes(pid)) + except (FileNotFoundError, ProcessLookupError, PermissionError): + # os.listdir() is gonna raise a lot of access denied + # exceptions in case of unprivileged user; that's fine + # as we'll just end up returning a connection with PID + # and fd set to None anyway. + # Both netstat -an and lsof does the same so it's + # unlikely we can do any better. + # ENOENT just means a PID disappeared on us. + continue + return inodes + + @staticmethod + def decode_address(addr, family): + """Accept an "ip:port" address as displayed in /proc/net/* + and convert it into a human readable form, like: + + "0500000A:0016" -> ("10.0.0.5", 22) + "0000000000000000FFFF00000100007F:9E49" -> ("::ffff:127.0.0.1", 40521) + + The IP address portion is a little or big endian four-byte + hexadecimal number; that is, the least significant byte is listed + first, so we need to reverse the order of the bytes to convert it + to an IP address. + The port is represented as a two-byte hexadecimal number. + + Reference: + http://linuxdevcenter.com/pub/a/linux/2000/11/16/LinuxAdmin.html + """ + ip, port = addr.split(':') + port = int(port, 16) + # this usually refers to a local socket in listen mode with + # no end-points connected + if not port: + return () + ip = ip.encode('ascii') + if family == socket.AF_INET: + # see: https://github.com/giampaolo/psutil/issues/201 + if LITTLE_ENDIAN: + ip = socket.inet_ntop(family, base64.b16decode(ip)[::-1]) + else: + ip = socket.inet_ntop(family, base64.b16decode(ip)) + else: # IPv6 + ip = base64.b16decode(ip) + try: + # see: https://github.com/giampaolo/psutil/issues/201 + if LITTLE_ENDIAN: + ip = socket.inet_ntop( + socket.AF_INET6, + struct.pack('>4I', *struct.unpack('<4I', ip)), + ) + else: + ip = socket.inet_ntop( + socket.AF_INET6, + struct.pack('<4I', *struct.unpack('<4I', ip)), + ) + except ValueError: + # see: https://github.com/giampaolo/psutil/issues/623 + if not supports_ipv6(): + raise _Ipv6UnsupportedError from None + raise + return ntp.addr(ip, port) + + @staticmethod + def process_inet(file, family, type_, inodes, filter_pid=None): + """Parse /proc/net/tcp* and /proc/net/udp* files.""" + if file.endswith('6') and not os.path.exists(file): + # IPv6 not supported + return + with open_text(file) as f: + f.readline() # skip the first line + for lineno, line in enumerate(f, 1): + try: + _, laddr, raddr, status, _, _, _, _, _, inode = ( + line.split()[:10] + ) + except ValueError: + msg = ( + f"error while parsing {file}; malformed line" + f" {lineno} {line!r}" + ) + raise RuntimeError(msg) from None + if inode in inodes: + # # We assume inet sockets are unique, so we error + # # out if there are multiple references to the + # # same inode. We won't do this for UNIX sockets. + # if len(inodes[inode]) > 1 and family != socket.AF_UNIX: + # raise ValueError("ambiguous inode with multiple " + # "PIDs references") + pid, fd = inodes[inode][0] + else: + pid, fd = None, -1 + if filter_pid is not None and filter_pid != pid: + continue + else: + if type_ == socket.SOCK_STREAM: + status = TCP_STATUSES[status] + else: + status = _common.CONN_NONE + try: + laddr = NetConnections.decode_address(laddr, family) + raddr = NetConnections.decode_address(raddr, family) + except _Ipv6UnsupportedError: + continue + yield (fd, family, type_, laddr, raddr, status, pid) + + @staticmethod + def process_unix(file, family, inodes, filter_pid=None): + """Parse /proc/net/unix files.""" + with open_text(file) as f: + f.readline() # skip the first line + for line in f: + tokens = line.split() + try: + _, _, _, _, type_, _, inode = tokens[0:7] + except ValueError: + if ' ' not in line: + # see: https://github.com/giampaolo/psutil/issues/766 + continue + msg = ( + f"error while parsing {file}; malformed line {line!r}" + ) + raise RuntimeError(msg) # noqa: B904 + if inode in inodes: # noqa: SIM108 + # With UNIX sockets we can have a single inode + # referencing many file descriptors. + pairs = inodes[inode] + else: + pairs = [(None, -1)] + for pid, fd in pairs: + if filter_pid is not None and filter_pid != pid: + continue + else: + path = tokens[-1] if len(tokens) == 8 else '' + type_ = _common.socktype_to_enum(int(type_)) + # XXX: determining the remote endpoint of a + # UNIX socket on Linux is not possible, see: + # https://serverfault.com/questions/252723/ + raddr = "" + status = _common.CONN_NONE + yield (fd, family, type_, path, raddr, status, pid) + + def retrieve(self, kind, pid=None): + self._procfs_path = get_procfs_path() + if pid is not None: + inodes = self.get_proc_inodes(pid) + if not inodes: + # no connections for this process + return [] + else: + inodes = self.get_all_inodes() + ret = set() + for proto_name, family, type_ in self.tmap[kind]: + path = f"{self._procfs_path}/net/{proto_name}" + if family in {socket.AF_INET, socket.AF_INET6}: + ls = self.process_inet( + path, family, type_, inodes, filter_pid=pid + ) + else: + ls = self.process_unix(path, family, inodes, filter_pid=pid) + for fd, family, type_, laddr, raddr, status, bound_pid in ls: + if pid: + conn = ntp.pconn(fd, family, type_, laddr, raddr, status) + else: + conn = ntp.sconn( + fd, family, type_, laddr, raddr, status, bound_pid + ) + ret.add(conn) + return list(ret) + + +_net_connections = NetConnections() + + +def net_connections(kind='inet'): + """Return system-wide open connections.""" + return _net_connections.retrieve(kind) + + +def net_io_counters(): + """Return network I/O statistics for every network interface + installed on the system as a dict of raw tuples. + """ + with open_text(f"{get_procfs_path()}/net/dev") as f: + lines = f.readlines() + retdict = {} + for line in lines[2:]: + colon = line.rfind(':') + assert colon > 0, repr(line) + name = line[:colon].strip() + fields = line[colon + 1 :].strip().split() + + ( + # in + bytes_recv, + packets_recv, + errin, + dropin, + _fifoin, # unused + _framein, # unused + _compressedin, # unused + _multicastin, # unused + # out + bytes_sent, + packets_sent, + errout, + dropout, + _fifoout, # unused + _collisionsout, # unused + _carrierout, # unused + _compressedout, # unused + ) = map(int, fields) + + retdict[name] = ( + bytes_sent, + bytes_recv, + packets_sent, + packets_recv, + errin, + errout, + dropin, + dropout, + ) + return retdict + + +def net_if_stats(): + """Get NIC stats (isup, duplex, speed, mtu).""" + duplex_map = { + cext.DUPLEX_FULL: NIC_DUPLEX_FULL, + cext.DUPLEX_HALF: NIC_DUPLEX_HALF, + cext.DUPLEX_UNKNOWN: NIC_DUPLEX_UNKNOWN, + } + names = net_io_counters().keys() + ret = {} + for name in names: + try: + mtu = cext.net_if_mtu(name) + flags = cext.net_if_flags(name) + duplex, speed = cext.net_if_duplex_speed(name) + except OSError as err: + # https://github.com/giampaolo/psutil/issues/1279 + if err.errno != errno.ENODEV: + raise + debug(err) + else: + output_flags = ','.join(flags) + isup = 'running' in flags + ret[name] = ntp.snicstats( + isup, duplex_map[duplex], speed, mtu, output_flags + ) + return ret + + +# ===================================================================== +# --- disks +# ===================================================================== + + +disk_usage = _psposix.disk_usage + + +def disk_io_counters(perdisk=False): + """Return disk I/O statistics for every disk installed on the + system as a dict of raw tuples. + """ + + def read_procfs(): + # OK, this is a bit confusing. The format of /proc/diskstats can + # have 3 variations. + # On Linux 2.4 each line has always 15 fields, e.g.: + # "3 0 8 hda 8 8 8 8 8 8 8 8 8 8 8" + # On Linux 2.6+ each line *usually* has 14 fields, and the disk + # name is in another position, like this: + # "3 0 hda 8 8 8 8 8 8 8 8 8 8 8" + # ...unless (Linux 2.6) the line refers to a partition instead + # of a disk, in which case the line has less fields (7): + # "3 1 hda1 8 8 8 8" + # 4.18+ has 4 fields added: + # "3 0 hda 8 8 8 8 8 8 8 8 8 8 8 0 0 0 0" + # 5.5 has 2 more fields. + # See: + # https://www.kernel.org/doc/Documentation/iostats.txt + # https://www.kernel.org/doc/Documentation/ABI/testing/procfs-diskstats + with open_text(f"{get_procfs_path()}/diskstats") as f: + lines = f.readlines() + for line in lines: + fields = line.split() + flen = len(fields) + # fmt: off + if flen == 15: + # Linux 2.4 + name = fields[3] + reads = int(fields[2]) + (reads_merged, rbytes, rtime, writes, writes_merged, + wbytes, wtime, _, busy_time, _) = map(int, fields[4:14]) + elif flen == 14 or flen >= 18: + # Linux 2.6+, line referring to a disk + name = fields[2] + (reads, reads_merged, rbytes, rtime, writes, writes_merged, + wbytes, wtime, _, busy_time, _) = map(int, fields[3:14]) + elif flen == 7: + # Linux 2.6+, line referring to a partition + name = fields[2] + reads, rbytes, writes, wbytes = map(int, fields[3:]) + rtime = wtime = reads_merged = writes_merged = busy_time = 0 + else: + msg = f"not sure how to interpret line {line!r}" + raise ValueError(msg) + yield (name, reads, writes, rbytes, wbytes, rtime, wtime, + reads_merged, writes_merged, busy_time) + # fmt: on + + def read_sysfs(): + for block in os.listdir('/sys/block'): + for root, _, files in os.walk(os.path.join('/sys/block', block)): + if 'stat' not in files: + continue + with open_text(os.path.join(root, 'stat')) as f: + fields = f.read().strip().split() + name = os.path.basename(root) + # fmt: off + (reads, reads_merged, rbytes, rtime, writes, writes_merged, + wbytes, wtime, _, busy_time) = map(int, fields[:10]) + yield (name, reads, writes, rbytes, wbytes, rtime, + wtime, reads_merged, writes_merged, busy_time) + # fmt: on + + if os.path.exists(f"{get_procfs_path()}/diskstats"): + gen = read_procfs() + elif os.path.exists('/sys/block'): + gen = read_sysfs() + else: + msg = ( + f"{get_procfs_path()}/diskstats nor /sys/block are available on" + " this system" + ) + raise NotImplementedError(msg) + + retdict = {} + for entry in gen: + # fmt: off + (name, reads, writes, rbytes, wbytes, rtime, wtime, reads_merged, + writes_merged, busy_time) = entry + if not perdisk and not is_storage_device(name): + # perdisk=False means we want to calculate totals so we skip + # partitions (e.g. 'sda1', 'nvme0n1p1') and only include + # base disk devices (e.g. 'sda', 'nvme0n1'). Base disks + # include a total of all their partitions + some extra size + # of their own: + # $ cat /proc/diskstats + # 259 0 sda 10485760 ... + # 259 1 sda1 5186039 ... + # 259 1 sda2 5082039 ... + # See: + # https://github.com/giampaolo/psutil/pull/1313 + continue + + rbytes *= DISK_SECTOR_SIZE + wbytes *= DISK_SECTOR_SIZE + retdict[name] = (reads, writes, rbytes, wbytes, rtime, wtime, + reads_merged, writes_merged, busy_time) + # fmt: on + + return retdict + + +class RootFsDeviceFinder: + """disk_partitions() may return partitions with device == "/dev/root" + or "rootfs". This container class uses different strategies to try to + obtain the real device path. Resources: + https://bootlin.com/blog/find-root-device/ + https://www.systutorials.com/how-to-find-the-disk-where-root-is-on-in-bash-on-linux/. + """ + + __slots__ = ['major', 'minor'] + + def __init__(self): + dev = os.stat("/").st_dev + self.major = os.major(dev) + self.minor = os.minor(dev) + + def ask_proc_partitions(self): + with open_text(f"{get_procfs_path()}/partitions") as f: + for line in f.readlines()[2:]: + fields = line.split() + if len(fields) < 4: # just for extra safety + continue + major = int(fields[0]) if fields[0].isdigit() else None + minor = int(fields[1]) if fields[1].isdigit() else None + name = fields[3] + if major == self.major and minor == self.minor: + if name: # just for extra safety + return f"/dev/{name}" + + def ask_sys_dev_block(self): + path = f"/sys/dev/block/{self.major}:{self.minor}/uevent" + with open_text(path) as f: + for line in f: + if line.startswith("DEVNAME="): + name = line.strip().rpartition("DEVNAME=")[2] + if name: # just for extra safety + return f"/dev/{name}" + + def ask_sys_class_block(self): + needle = f"{self.major}:{self.minor}" + files = glob.iglob("/sys/class/block/*/dev") + for file in files: + try: + f = open_text(file) + except FileNotFoundError: # race condition + continue + else: + with f: + data = f.read().strip() + if data == needle: + name = os.path.basename(os.path.dirname(file)) + return f"/dev/{name}" + + def find(self): + path = None + if path is None: + try: + path = self.ask_proc_partitions() + except OSError as err: + debug(err) + if path is None: + try: + path = self.ask_sys_dev_block() + except OSError as err: + debug(err) + if path is None: + try: + path = self.ask_sys_class_block() + except OSError as err: + debug(err) + # We use exists() because the "/dev/*" part of the path is hard + # coded, so we want to be sure. + if path is not None and os.path.exists(path): + return path + + +def disk_partitions(all=False): + """Return mounted disk partitions as a list of namedtuples.""" + fstypes = set() + procfs_path = get_procfs_path() + if not all: + with open_text(f"{procfs_path}/filesystems") as f: + for line in f: + line = line.strip() + if not line.startswith("nodev"): + fstypes.add(line.strip()) + else: + # ignore all lines starting with "nodev" except "nodev zfs" + fstype = line.split("\t")[1] + if fstype == "zfs": + fstypes.add("zfs") + + # See: https://github.com/giampaolo/psutil/issues/1307 + if procfs_path == "/proc" and os.path.isfile('/etc/mtab'): + mounts_path = os.path.realpath("/etc/mtab") + else: + mounts_path = os.path.realpath(f"{procfs_path}/self/mounts") + + retlist = [] + partitions = cext.disk_partitions(mounts_path) + for partition in partitions: + device, mountpoint, fstype, opts = partition + if device == 'none': + device = '' + if device in {"/dev/root", "rootfs"}: + device = RootFsDeviceFinder().find() or device + if not all: + if not device or fstype not in fstypes: + continue + ntuple = ntp.sdiskpart(device, mountpoint, fstype, opts) + retlist.append(ntuple) + + return retlist + + +# ===================================================================== +# --- sensors +# ===================================================================== + + +def sensors_temperatures(): + """Return hardware (CPU and others) temperatures as a dict + including hardware name, label, current, max and critical + temperatures. + + Implementation notes: + - /sys/class/hwmon looks like the most recent interface to + retrieve this info, and this implementation relies on it + only (old distros will probably use something else) + - lm-sensors on Ubuntu 16.04 relies on /sys/class/hwmon + - /sys/class/thermal/thermal_zone* is another one but it's more + difficult to parse + """ + ret = collections.defaultdict(list) + basenames = glob.glob('/sys/class/hwmon/hwmon*/temp*_*') + # CentOS has an intermediate /device directory: + # https://github.com/giampaolo/psutil/issues/971 + # https://github.com/nicolargo/glances/issues/1060 + basenames.extend(glob.glob('/sys/class/hwmon/hwmon*/device/temp*_*')) + basenames = sorted({x.split('_')[0] for x in basenames}) + + # Only add the coretemp hwmon entries if they're not already in + # /sys/class/hwmon/ + # https://github.com/giampaolo/psutil/issues/1708 + # https://github.com/giampaolo/psutil/pull/1648 + basenames2 = glob.glob( + '/sys/devices/platform/coretemp.*/hwmon/hwmon*/temp*_*' + ) + repl = re.compile(r"/sys/devices/platform/coretemp.*/hwmon/") + for name in basenames2: + altname = repl.sub('/sys/class/hwmon/', name) + if altname not in basenames: + basenames.append(name) + + for base in basenames: + try: + path = base + '_input' + current = float(bcat(path)) / 1000.0 + path = os.path.join(os.path.dirname(base), 'name') + unit_name = cat(path).strip() + except (OSError, ValueError): + # A lot of things can go wrong here, so let's just skip the + # whole entry. Sure thing is Linux's /sys/class/hwmon really + # is a stinky broken mess. + # https://github.com/giampaolo/psutil/issues/1009 + # https://github.com/giampaolo/psutil/issues/1101 + # https://github.com/giampaolo/psutil/issues/1129 + # https://github.com/giampaolo/psutil/issues/1245 + # https://github.com/giampaolo/psutil/issues/1323 + continue + + high = bcat(base + '_max', fallback=None) + critical = bcat(base + '_crit', fallback=None) + label = cat(base + '_label', fallback='').strip() + + if high is not None: + try: + high = float(high) / 1000.0 + except ValueError: + high = None + if critical is not None: + try: + critical = float(critical) / 1000.0 + except ValueError: + critical = None + + ret[unit_name].append((label, current, high, critical)) + + # Indication that no sensors were detected in /sys/class/hwmon/ + if not basenames: + basenames = glob.glob('/sys/class/thermal/thermal_zone*') + basenames = sorted(set(basenames)) + + for base in basenames: + try: + path = os.path.join(base, 'temp') + current = float(bcat(path)) / 1000.0 + path = os.path.join(base, 'type') + unit_name = cat(path).strip() + except (OSError, ValueError) as err: + debug(err) + continue + + trip_paths = glob.glob(base + '/trip_point*') + trip_points = { + '_'.join(os.path.basename(p).split('_')[0:3]) + for p in trip_paths + } + critical = None + high = None + for trip_point in trip_points: + path = os.path.join(base, trip_point + "_type") + trip_type = cat(path, fallback='').strip() + if trip_type == 'critical': + critical = bcat( + os.path.join(base, trip_point + "_temp"), fallback=None + ) + elif trip_type == 'high': + high = bcat( + os.path.join(base, trip_point + "_temp"), fallback=None + ) + + if high is not None: + try: + high = float(high) / 1000.0 + except ValueError: + high = None + if critical is not None: + try: + critical = float(critical) / 1000.0 + except ValueError: + critical = None + + ret[unit_name].append(('', current, high, critical)) + + return dict(ret) + + +def sensors_fans(): + """Return hardware fans info (for CPU and other peripherals) as a + dict including hardware label and current speed. + + Implementation notes: + - /sys/class/hwmon looks like the most recent interface to + retrieve this info, and this implementation relies on it + only (old distros will probably use something else) + - lm-sensors on Ubuntu 16.04 relies on /sys/class/hwmon + """ + ret = collections.defaultdict(list) + basenames = glob.glob('/sys/class/hwmon/hwmon*/fan*_*') + if not basenames: + # CentOS has an intermediate /device directory: + # https://github.com/giampaolo/psutil/issues/971 + basenames = glob.glob('/sys/class/hwmon/hwmon*/device/fan*_*') + + basenames = sorted({x.split("_")[0] for x in basenames}) + for base in basenames: + try: + current = int(bcat(base + '_input')) + except OSError as err: + debug(err) + continue + unit_name = cat(os.path.join(os.path.dirname(base), 'name')).strip() + label = cat(base + '_label', fallback='').strip() + ret[unit_name].append(ntp.sfan(label, current)) + + return dict(ret) + + +def sensors_battery(): + """Return battery information. + Implementation note: it appears /sys/class/power_supply/BAT0/ + directory structure may vary and provide files with the same + meaning but under different names, see: + https://github.com/giampaolo/psutil/issues/966. + """ + null = object() + + def multi_bcat(*paths): + """Attempt to read the content of multiple files which may + not exist. If none of them exist return None. + """ + for path in paths: + ret = bcat(path, fallback=null) + if ret != null: + try: + return int(ret) + except ValueError: + return ret.strip() + return None + + bats = [ + x + for x in os.listdir(POWER_SUPPLY_PATH) + if x.startswith('BAT') or 'battery' in x.lower() + ] + if not bats: + return None + # Get the first available battery. Usually this is "BAT0", except + # some rare exceptions: + # https://github.com/giampaolo/psutil/issues/1238 + root = os.path.join(POWER_SUPPLY_PATH, min(bats)) + + # Base metrics. + energy_now = multi_bcat(root + "/energy_now", root + "/charge_now") + power_now = multi_bcat(root + "/power_now", root + "/current_now") + energy_full = multi_bcat(root + "/energy_full", root + "/charge_full") + time_to_empty = multi_bcat(root + "/time_to_empty_now") + + # Percent. If we have energy_full the percentage will be more + # accurate compared to reading /capacity file (float vs. int). + if energy_full is not None and energy_now is not None: + try: + percent = 100.0 * energy_now / energy_full + except ZeroDivisionError: + percent = 0.0 + else: + percent = int(cat(root + "/capacity", fallback=-1)) + if percent == -1: + return None + + # Is AC power cable plugged in? + # Note: AC0 is not always available and sometimes (e.g. CentOS7) + # it's called "AC". + power_plugged = None + online = multi_bcat( + os.path.join(POWER_SUPPLY_PATH, "AC0/online"), + os.path.join(POWER_SUPPLY_PATH, "AC/online"), + ) + if online is not None: + power_plugged = online == 1 + else: + status = cat(root + "/status", fallback="").strip().lower() + if status == "discharging": + power_plugged = False + elif status in {"charging", "full"}: + power_plugged = True + + # Seconds left. + # Note to self: we may also calculate the charging ETA as per: + # https://github.com/thialfihar/dotfiles/blob/ + # 013937745fd9050c30146290e8f963d65c0179e6/bin/battery.py#L55 + if power_plugged: + secsleft = _common.POWER_TIME_UNLIMITED + elif energy_now is not None and power_now is not None: + try: + secsleft = int(energy_now / abs(power_now) * 3600) + except ZeroDivisionError: + secsleft = _common.POWER_TIME_UNKNOWN + elif time_to_empty is not None: + secsleft = int(time_to_empty * 60) + if secsleft < 0: + secsleft = _common.POWER_TIME_UNKNOWN + else: + secsleft = _common.POWER_TIME_UNKNOWN + + return ntp.sbattery(percent, secsleft, power_plugged) + + +# ===================================================================== +# --- other system functions +# ===================================================================== + + +def users(): + """Return currently connected users as a list of namedtuples.""" + retlist = [] + rawlist = cext.users() + for item in rawlist: + user, tty, hostname, tstamp, pid = item + nt = ntp.suser(user, tty or None, hostname, tstamp, pid) + retlist.append(nt) + return retlist + + +def boot_time(): + """Return the system boot time expressed in seconds since the epoch.""" + path = f"{get_procfs_path()}/stat" + with open_binary(path) as f: + for line in f: + if line.startswith(b'btime'): + return float(line.strip().split()[1]) + msg = f"line 'btime' not found in {path}" + raise RuntimeError(msg) + + +# ===================================================================== +# --- processes +# ===================================================================== + + +def pids(): + """Returns a list of PIDs currently running on the system.""" + path = get_procfs_path().encode(ENCODING) + return [int(x) for x in os.listdir(path) if x.isdigit()] + + +def pid_exists(pid): + """Check for the existence of a unix PID. Linux TIDs are not + supported (always return False). + """ + if not _psposix.pid_exists(pid): + return False + else: + # Linux's apparently does not distinguish between PIDs and TIDs + # (thread IDs). + # listdir("/proc") won't show any TID (only PIDs) but + # os.stat("/proc/{tid}") will succeed if {tid} exists. + # os.kill() can also be passed a TID. This is quite confusing. + # In here we want to enforce this distinction and support PIDs + # only, see: + # https://github.com/giampaolo/psutil/issues/687 + try: + # Note: already checked that this is faster than using a + # regular expr. Also (a lot) faster than doing + # 'return pid in pids()' + path = f"{get_procfs_path()}/{pid}/status" + with open_binary(path) as f: + for line in f: + if line.startswith(b"Tgid:"): + tgid = int(line.split()[1]) + # If tgid and pid are the same then we're + # dealing with a process PID. + return tgid == pid + msg = f"'Tgid' line not found in {path}" + raise ValueError(msg) + except (OSError, ValueError): + return pid in pids() + + +def ppid_map(): + """Obtain a {pid: ppid, ...} dict for all running processes in + one shot. Used to speed up Process.children(). + """ + ret = {} + procfs_path = get_procfs_path() + for pid in pids(): + try: + with open_binary(f"{procfs_path}/{pid}/stat") as f: + data = f.read() + except (FileNotFoundError, ProcessLookupError): + pass + except PermissionError as err: + raise AccessDenied(pid) from err + else: + rpar = data.rfind(b')') + dset = data[rpar + 2 :].split() + ppid = int(dset[1]) + ret[pid] = ppid + return ret + + +def wrap_exceptions(fun): + """Decorator which translates bare OSError exceptions into + NoSuchProcess and AccessDenied. + """ + + @functools.wraps(fun) + def wrapper(self, *args, **kwargs): + pid, name = self.pid, self._name + try: + return fun(self, *args, **kwargs) + except PermissionError as err: + raise AccessDenied(pid, name) from err + except ProcessLookupError as err: + self._raise_if_zombie() + raise NoSuchProcess(pid, name) from err + except FileNotFoundError as err: + self._raise_if_zombie() + # /proc/PID directory may still exist, but the files within + # it may not, indicating the process is gone, see: + # https://github.com/giampaolo/psutil/issues/2418 + if not os.path.exists(f"{self._procfs_path}/{pid}/stat"): + raise NoSuchProcess(pid, name) from err + raise + + return wrapper + + +class Process: + """Linux process implementation.""" + + __slots__ = [ + "_cache", + "_ctime", + "_name", + "_ppid", + "_procfs_path", + "pid", + ] + + def __init__(self, pid): + self.pid = pid + self._name = None + self._ppid = None + self._ctime = None + self._procfs_path = get_procfs_path() + + def _is_zombie(self): + # Note: most of the times Linux is able to return info about the + # process even if it's a zombie, and /proc/{pid} will exist. + # There are some exceptions though, like exe(), cmdline() and + # memory_maps(). In these cases /proc/{pid}/{file} exists but + # it's empty. Instead of returning a "null" value we'll raise an + # exception. + try: + data = bcat(f"{self._procfs_path}/{self.pid}/stat") + except OSError: + return False + else: + rpar = data.rfind(b')') + status = data[rpar + 2 : rpar + 3] + return status == b"Z" + + def _raise_if_zombie(self): + if self._is_zombie(): + raise ZombieProcess(self.pid, self._name, self._ppid) + + def _raise_if_not_alive(self): + """Raise NSP if the process disappeared on us.""" + # For those C function who do not raise NSP, possibly returning + # incorrect or incomplete result. + os.stat(f"{self._procfs_path}/{self.pid}") + + def _readlink(self, path, fallback=UNSET): + # * https://github.com/giampaolo/psutil/issues/503 + # os.readlink('/proc/pid/exe') may raise ESRCH (ProcessLookupError) + # instead of ENOENT (FileNotFoundError) when it races. + # * ENOENT may occur also if the path actually exists if PID is + # a low PID (~0-20 range). + # * https://github.com/giampaolo/psutil/issues/2514 + try: + return readlink(path) + except (FileNotFoundError, ProcessLookupError): + if os.path.lexists(f"{self._procfs_path}/{self.pid}"): + self._raise_if_zombie() + if fallback is not UNSET: + return fallback + raise + + @wrap_exceptions + @memoize_when_activated + def _parse_stat_file(self): + """Parse /proc/{pid}/stat file and return a dict with various + process info. + Using "man proc" as a reference: where "man proc" refers to + position N always subtract 3 (e.g ppid position 4 in + 'man proc' == position 1 in here). + The return value is cached in case oneshot() ctx manager is + in use. + """ + data = bcat(f"{self._procfs_path}/{self.pid}/stat") + # Process name is between parentheses. It can contain spaces and + # other parentheses. This is taken into account by looking for + # the first occurrence of "(" and the last occurrence of ")". + rpar = data.rfind(b')') + name = data[data.find(b'(') + 1 : rpar] + fields = data[rpar + 2 :].split() + + ret = {} + ret['name'] = name + ret['status'] = fields[0] + ret['ppid'] = fields[1] + ret['ttynr'] = fields[4] + ret['utime'] = fields[11] + ret['stime'] = fields[12] + ret['children_utime'] = fields[13] + ret['children_stime'] = fields[14] + ret['create_time'] = fields[19] + ret['cpu_num'] = fields[36] + try: + ret['blkio_ticks'] = fields[39] # aka 'delayacct_blkio_ticks' + except IndexError: + # https://github.com/giampaolo/psutil/issues/2455 + debug("can't get blkio_ticks, set iowait to 0") + ret['blkio_ticks'] = 0 + + return ret + + @wrap_exceptions + @memoize_when_activated + def _read_status_file(self): + """Read /proc/{pid}/stat file and return its content. + The return value is cached in case oneshot() ctx manager is + in use. + """ + with open_binary(f"{self._procfs_path}/{self.pid}/status") as f: + return f.read() + + @wrap_exceptions + @memoize_when_activated + def _read_smaps_file(self): + with open_binary(f"{self._procfs_path}/{self.pid}/smaps") as f: + return f.read().strip() + + def oneshot_enter(self): + self._parse_stat_file.cache_activate(self) + self._read_status_file.cache_activate(self) + self._read_smaps_file.cache_activate(self) + + def oneshot_exit(self): + self._parse_stat_file.cache_deactivate(self) + self._read_status_file.cache_deactivate(self) + self._read_smaps_file.cache_deactivate(self) + + @wrap_exceptions + def name(self): + # XXX - gets changed later and probably needs refactoring + return decode(self._parse_stat_file()['name']) + + @wrap_exceptions + def exe(self): + return self._readlink( + f"{self._procfs_path}/{self.pid}/exe", fallback="" + ) + + @wrap_exceptions + def cmdline(self): + with open_text(f"{self._procfs_path}/{self.pid}/cmdline") as f: + data = f.read() + if not data: + # may happen in case of zombie process + self._raise_if_zombie() + return [] + # 'man proc' states that args are separated by null bytes '\0' + # and last char is supposed to be a null byte. Nevertheless + # some processes may change their cmdline after being started + # (via setproctitle() or similar), they are usually not + # compliant with this rule and use spaces instead. Google + # Chrome process is an example. See: + # https://github.com/giampaolo/psutil/issues/1179 + sep = '\x00' if data.endswith('\x00') else ' ' + if data.endswith(sep): + data = data[:-1] + cmdline = data.split(sep) + # Sometimes last char is a null byte '\0' but the args are + # separated by spaces, see: https://github.com/giampaolo/psutil/ + # issues/1179#issuecomment-552984549 + if sep == '\x00' and len(cmdline) == 1 and ' ' in data: + cmdline = data.split(' ') + return cmdline + + @wrap_exceptions + def environ(self): + with open_text(f"{self._procfs_path}/{self.pid}/environ") as f: + data = f.read() + return parse_environ_block(data) + + @wrap_exceptions + def terminal(self): + tty_nr = int(self._parse_stat_file()['ttynr']) + tmap = _psposix.get_terminal_map() + try: + return tmap[tty_nr] + except KeyError: + return None + + # May not be available on old kernels. + if os.path.exists(f"/proc/{os.getpid()}/io"): + + @wrap_exceptions + def io_counters(self): + fname = f"{self._procfs_path}/{self.pid}/io" + fields = {} + with open_binary(fname) as f: + for line in f: + # https://github.com/giampaolo/psutil/issues/1004 + line = line.strip() + if line: + try: + name, value = line.split(b': ') + except ValueError: + # https://github.com/giampaolo/psutil/issues/1004 + continue + else: + fields[name] = int(value) + if not fields: + msg = f"{fname} file was empty" + raise RuntimeError(msg) + try: + return ntp.pio( + fields[b'syscr'], # read syscalls + fields[b'syscw'], # write syscalls + fields[b'read_bytes'], # read bytes + fields[b'write_bytes'], # write bytes + fields[b'rchar'], # read chars + fields[b'wchar'], # write chars + ) + except KeyError as err: + msg = ( + f"{err.args[0]!r} field was not found in {fname}; found" + f" fields are {fields!r}" + ) + raise ValueError(msg) from None + + @wrap_exceptions + def cpu_times(self): + values = self._parse_stat_file() + utime = float(values['utime']) / CLOCK_TICKS + stime = float(values['stime']) / CLOCK_TICKS + children_utime = float(values['children_utime']) / CLOCK_TICKS + children_stime = float(values['children_stime']) / CLOCK_TICKS + iowait = float(values['blkio_ticks']) / CLOCK_TICKS + return ntp.pcputimes( + utime, stime, children_utime, children_stime, iowait + ) + + @wrap_exceptions + def cpu_num(self): + """What CPU the process is on.""" + return int(self._parse_stat_file()['cpu_num']) + + @wrap_exceptions + def wait(self, timeout=None): + return _psposix.wait_pid(self.pid, timeout) + + @wrap_exceptions + def create_time(self, monotonic=False): + # The 'starttime' field in /proc/[pid]/stat is expressed in + # jiffies (clock ticks per second), a relative value which + # represents the number of clock ticks that have passed since + # the system booted until the process was created. It never + # changes and is unaffected by system clock updates. + if self._ctime is None: + self._ctime = ( + float(self._parse_stat_file()['create_time']) / CLOCK_TICKS + ) + if monotonic: + return self._ctime + # Add the boot time, returning time expressed in seconds since + # the epoch. This is subject to system clock updates. + return self._ctime + boot_time() + + @wrap_exceptions + def memory_info(self): + # ============================================================ + # | FIELD | DESCRIPTION | AKA | TOP | + # ============================================================ + # | rss | resident set size | | RES | + # | vms | total program size | size | VIRT | + # | shared | shared pages (from shared mappings) | | SHR | + # | text | text ('code') | trs | CODE | + # | lib | library (unused in Linux 2.6) | lrs | | + # | data | data + stack | drs | DATA | + # | dirty | dirty pages (unused in Linux 2.6) | dt | | + # ============================================================ + with open_binary(f"{self._procfs_path}/{self.pid}/statm") as f: + vms, rss, shared, text, lib, data, dirty = ( + int(x) * PAGESIZE for x in f.readline().split()[:7] + ) + return ntp.pmem(rss, vms, shared, text, lib, data, dirty) + + if HAS_PROC_SMAPS_ROLLUP or HAS_PROC_SMAPS: + + def _parse_smaps_rollup(self): + # /proc/pid/smaps_rollup was added to Linux in 2017. Faster + # than /proc/pid/smaps. It reports higher PSS than */smaps + # (from 1k up to 200k higher; tested against all processes). + # IMPORTANT: /proc/pid/smaps_rollup is weird, because it + # raises ESRCH / ENOENT for many PIDs, even if they're alive + # (also as root). In that case we'll use /proc/pid/smaps as + # fallback, which is slower but has a +50% success rate + # compared to /proc/pid/smaps_rollup. + uss = pss = swap = 0 + with open_binary( + f"{self._procfs_path}/{self.pid}/smaps_rollup" + ) as f: + for line in f: + if line.startswith(b"Private_"): + # Private_Clean, Private_Dirty, Private_Hugetlb + uss += int(line.split()[1]) * 1024 + elif line.startswith(b"Pss:"): + pss = int(line.split()[1]) * 1024 + elif line.startswith(b"Swap:"): + swap = int(line.split()[1]) * 1024 + return (uss, pss, swap) + + @wrap_exceptions + def _parse_smaps( + self, + # Gets Private_Clean, Private_Dirty, Private_Hugetlb. + _private_re=re.compile(br"\nPrivate.*:\s+(\d+)"), + _pss_re=re.compile(br"\nPss\:\s+(\d+)"), + _swap_re=re.compile(br"\nSwap\:\s+(\d+)"), + ): + # /proc/pid/smaps does not exist on kernels < 2.6.14 or if + # CONFIG_MMU kernel configuration option is not enabled. + + # Note: using 3 regexes is faster than reading the file + # line by line. + # + # You might be tempted to calculate USS by subtracting + # the "shared" value from the "resident" value in + # /proc//statm. But at least on Linux, statm's "shared" + # value actually counts pages backed by files, which has + # little to do with whether the pages are actually shared. + # /proc/self/smaps on the other hand appears to give us the + # correct information. + smaps_data = self._read_smaps_file() + # Note: smaps file can be empty for certain processes. + # The code below will not crash though and will result to 0. + uss = sum(map(int, _private_re.findall(smaps_data))) * 1024 + pss = sum(map(int, _pss_re.findall(smaps_data))) * 1024 + swap = sum(map(int, _swap_re.findall(smaps_data))) * 1024 + return (uss, pss, swap) + + @wrap_exceptions + def memory_full_info(self): + if HAS_PROC_SMAPS_ROLLUP: # faster + try: + uss, pss, swap = self._parse_smaps_rollup() + except (ProcessLookupError, FileNotFoundError): + uss, pss, swap = self._parse_smaps() + else: + uss, pss, swap = self._parse_smaps() + basic_mem = self.memory_info() + return ntp.pfullmem(*basic_mem + (uss, pss, swap)) + + else: + memory_full_info = memory_info + + if HAS_PROC_SMAPS: + + @wrap_exceptions + def memory_maps(self): + """Return process's mapped memory regions as a list of named + tuples. Fields are explained in 'man proc'; here is an updated + (Apr 2012) version: https://git.kernel.org/pub/scm/linux/kernel/git/torvalds/linux.git/tree/Documentation/filesystems/proc.txt?id=b76437579d1344b612cf1851ae610c636cec7db0. + + /proc/{PID}/smaps does not exist on kernels < 2.6.14 or if + CONFIG_MMU kernel configuration option is not enabled. + """ + + def get_blocks(lines, current_block): + data = {} + for line in lines: + fields = line.split(None, 5) + if not fields[0].endswith(b':'): + # new block section + yield (current_block.pop(), data) + current_block.append(line) + else: + try: + data[fields[0]] = int(fields[1]) * 1024 + except (ValueError, IndexError): + if fields[0].startswith(b'VmFlags:'): + # see issue #369 + continue + msg = f"don't know how to interpret line {line!r}" + raise ValueError(msg) from None + yield (current_block.pop(), data) + + data = self._read_smaps_file() + # Note: smaps file can be empty for certain processes or for + # zombies. + if not data: + self._raise_if_zombie() + return [] + lines = data.split(b'\n') + ls = [] + first_line = lines.pop(0) + current_block = [first_line] + for header, data in get_blocks(lines, current_block): + hfields = header.split(None, 5) + try: + addr, perms, _offset, _dev, _inode, path = hfields + except ValueError: + addr, perms, _offset, _dev, _inode, path = hfields + [''] + if not path: + path = '[anon]' + else: + path = decode(path) + path = path.strip() + if path.endswith(' (deleted)') and not path_exists_strict( + path + ): + path = path[:-10] + item = ( + decode(addr), + decode(perms), + path, + data.get(b'Rss:', 0), + data.get(b'Size:', 0), + data.get(b'Pss:', 0), + data.get(b'Shared_Clean:', 0), + data.get(b'Shared_Dirty:', 0), + data.get(b'Private_Clean:', 0), + data.get(b'Private_Dirty:', 0), + data.get(b'Referenced:', 0), + data.get(b'Anonymous:', 0), + data.get(b'Swap:', 0), + ) + ls.append(item) + return ls + + @wrap_exceptions + def cwd(self): + return self._readlink( + f"{self._procfs_path}/{self.pid}/cwd", fallback="" + ) + + @wrap_exceptions + def num_ctx_switches( + self, _ctxsw_re=re.compile(br'ctxt_switches:\t(\d+)') + ): + data = self._read_status_file() + ctxsw = _ctxsw_re.findall(data) + if not ctxsw: + msg = ( + "'voluntary_ctxt_switches' and" + " 'nonvoluntary_ctxt_switches'lines were not found in" + f" {self._procfs_path}/{self.pid}/status; the kernel is" + " probably older than 2.6.23" + ) + raise NotImplementedError(msg) + return ntp.pctxsw(int(ctxsw[0]), int(ctxsw[1])) + + @wrap_exceptions + def num_threads(self, _num_threads_re=re.compile(br'Threads:\t(\d+)')): + # Using a re is faster than iterating over file line by line. + data = self._read_status_file() + return int(_num_threads_re.findall(data)[0]) + + @wrap_exceptions + def threads(self): + thread_ids = os.listdir(f"{self._procfs_path}/{self.pid}/task") + thread_ids.sort() + retlist = [] + hit_enoent = False + for thread_id in thread_ids: + fname = f"{self._procfs_path}/{self.pid}/task/{thread_id}/stat" + try: + with open_binary(fname) as f: + st = f.read().strip() + except (FileNotFoundError, ProcessLookupError): + # no such file or directory or no such process; + # it means thread disappeared on us + hit_enoent = True + continue + # ignore the first two values ("pid (exe)") + st = st[st.find(b')') + 2 :] + values = st.split(b' ') + utime = float(values[11]) / CLOCK_TICKS + stime = float(values[12]) / CLOCK_TICKS + ntuple = ntp.pthread(int(thread_id), utime, stime) + retlist.append(ntuple) + if hit_enoent: + self._raise_if_not_alive() + return retlist + + @wrap_exceptions + def nice_get(self): + # with open_text(f"{self._procfs_path}/{self.pid}/stat") as f: + # data = f.read() + # return int(data.split()[18]) + + # Use C implementation + return cext.proc_priority_get(self.pid) + + @wrap_exceptions + def nice_set(self, value): + return cext.proc_priority_set(self.pid, value) + + # starting from CentOS 6. + if HAS_CPU_AFFINITY: + + @wrap_exceptions + def cpu_affinity_get(self): + return cext.proc_cpu_affinity_get(self.pid) + + def _get_eligible_cpus( + self, _re=re.compile(br"Cpus_allowed_list:\t(\d+)-(\d+)") + ): + # See: https://github.com/giampaolo/psutil/issues/956 + data = self._read_status_file() + match = _re.findall(data) + if match: + return list(range(int(match[0][0]), int(match[0][1]) + 1)) + else: + return list(range(len(per_cpu_times()))) + + @wrap_exceptions + def cpu_affinity_set(self, cpus): + try: + cext.proc_cpu_affinity_set(self.pid, cpus) + except (OSError, ValueError) as err: + if isinstance(err, ValueError) or err.errno == errno.EINVAL: + eligible_cpus = self._get_eligible_cpus() + all_cpus = tuple(range(len(per_cpu_times()))) + for cpu in cpus: + if cpu not in all_cpus: + msg = ( + f"invalid CPU {cpu!r}; choose between" + f" {eligible_cpus!r}" + ) + raise ValueError(msg) from None + if cpu not in eligible_cpus: + msg = ( + f"CPU number {cpu} is not eligible; choose" + f" between {eligible_cpus}" + ) + raise ValueError(msg) from err + raise + + # only starting from kernel 2.6.13 + if HAS_PROC_IO_PRIORITY: + + @wrap_exceptions + def ionice_get(self): + ioclass, value = cext.proc_ioprio_get(self.pid) + ioclass = IOPriority(ioclass) + return ntp.pionice(ioclass, value) + + @wrap_exceptions + def ionice_set(self, ioclass, value): + if value is None: + value = 0 + if value and ioclass in { + IOPriority.IOPRIO_CLASS_IDLE, + IOPriority.IOPRIO_CLASS_NONE, + }: + msg = f"{ioclass!r} ioclass accepts no value" + raise ValueError(msg) + if value < 0 or value > 7: + msg = "value not in 0-7 range" + raise ValueError(msg) + return cext.proc_ioprio_set(self.pid, ioclass, value) + + if hasattr(resource, "prlimit"): + + @wrap_exceptions + def rlimit(self, resource_, limits=None): + # If pid is 0 prlimit() applies to the calling process and + # we don't want that. We should never get here though as + # PID 0 is not supported on Linux. + if self.pid == 0: + msg = "can't use prlimit() against PID 0 process" + raise ValueError(msg) + try: + if limits is None: + # get + return resource.prlimit(self.pid, resource_) + else: + # set + if len(limits) != 2: + msg = ( + "second argument must be a (soft, hard) " + f"tuple, got {limits!r}" + ) + raise ValueError(msg) + resource.prlimit(self.pid, resource_, limits) + except OSError as err: + if err.errno == errno.ENOSYS: + # I saw this happening on Travis: + # https://travis-ci.org/giampaolo/psutil/jobs/51368273 + self._raise_if_zombie() + raise + + @wrap_exceptions + def status(self): + letter = self._parse_stat_file()['status'] + letter = letter.decode() + # XXX is '?' legit? (we're not supposed to return it anyway) + return PROC_STATUSES.get(letter, '?') + + @wrap_exceptions + def open_files(self): + retlist = [] + files = os.listdir(f"{self._procfs_path}/{self.pid}/fd") + hit_enoent = False + for fd in files: + file = f"{self._procfs_path}/{self.pid}/fd/{fd}" + try: + path = readlink(file) + except (FileNotFoundError, ProcessLookupError): + # ENOENT == file which is gone in the meantime + hit_enoent = True + continue + except OSError as err: + if err.errno == errno.EINVAL: + # not a link + continue + if err.errno == errno.ENAMETOOLONG: + # file name too long + debug(err) + continue + raise + else: + # If path is not an absolute there's no way to tell + # whether it's a regular file or not, so we skip it. + # A regular file is always supposed to be have an + # absolute path though. + if path.startswith('/') and isfile_strict(path): + # Get file position and flags. + file = f"{self._procfs_path}/{self.pid}/fdinfo/{fd}" + try: + with open_binary(file) as f: + pos = int(f.readline().split()[1]) + flags = int(f.readline().split()[1], 8) + except (FileNotFoundError, ProcessLookupError): + # fd gone in the meantime; process may + # still be alive + hit_enoent = True + else: + mode = file_flags_to_mode(flags) + ntuple = ntp.popenfile( + path, int(fd), int(pos), mode, flags + ) + retlist.append(ntuple) + if hit_enoent: + self._raise_if_not_alive() + return retlist + + @wrap_exceptions + def net_connections(self, kind='inet'): + ret = _net_connections.retrieve(kind, self.pid) + self._raise_if_not_alive() + return ret + + @wrap_exceptions + def num_fds(self): + return len(os.listdir(f"{self._procfs_path}/{self.pid}/fd")) + + @wrap_exceptions + def ppid(self): + return int(self._parse_stat_file()['ppid']) + + @wrap_exceptions + def uids(self, _uids_re=re.compile(br'Uid:\t(\d+)\t(\d+)\t(\d+)')): + data = self._read_status_file() + real, effective, saved = _uids_re.findall(data)[0] + return ntp.puids(int(real), int(effective), int(saved)) + + @wrap_exceptions + def gids(self, _gids_re=re.compile(br'Gid:\t(\d+)\t(\d+)\t(\d+)')): + data = self._read_status_file() + real, effective, saved = _gids_re.findall(data)[0] + return ntp.pgids(int(real), int(effective), int(saved)) diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/psutil/_psosx.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/psutil/_psosx.py new file mode 100644 index 0000000000000000000000000000000000000000..b14d5e023bec1f9cb5cb06b4caf7170489f3d4c7 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/psutil/_psosx.py @@ -0,0 +1,549 @@ +# Copyright (c) 2009, Giampaolo Rodola'. All rights reserved. +# Use of this source code is governed by a BSD-style license that can be +# found in the LICENSE file. + +"""macOS platform implementation.""" + +import errno +import functools +import os + +from . import _common +from . import _ntuples as ntp +from . import _psposix +from . import _psutil_osx as cext +from ._common import AccessDenied +from ._common import NoSuchProcess +from ._common import ZombieProcess +from ._common import conn_tmap +from ._common import conn_to_ntuple +from ._common import debug +from ._common import isfile_strict +from ._common import memoize_when_activated +from ._common import parse_environ_block +from ._common import usage_percent + +__extra__all__ = [] + + +# ===================================================================== +# --- globals +# ===================================================================== + + +PAGESIZE = cext.getpagesize() +AF_LINK = cext.AF_LINK + +TCP_STATUSES = { + cext.TCPS_ESTABLISHED: _common.CONN_ESTABLISHED, + cext.TCPS_SYN_SENT: _common.CONN_SYN_SENT, + cext.TCPS_SYN_RECEIVED: _common.CONN_SYN_RECV, + cext.TCPS_FIN_WAIT_1: _common.CONN_FIN_WAIT1, + cext.TCPS_FIN_WAIT_2: _common.CONN_FIN_WAIT2, + cext.TCPS_TIME_WAIT: _common.CONN_TIME_WAIT, + cext.TCPS_CLOSED: _common.CONN_CLOSE, + cext.TCPS_CLOSE_WAIT: _common.CONN_CLOSE_WAIT, + cext.TCPS_LAST_ACK: _common.CONN_LAST_ACK, + cext.TCPS_LISTEN: _common.CONN_LISTEN, + cext.TCPS_CLOSING: _common.CONN_CLOSING, + cext.PSUTIL_CONN_NONE: _common.CONN_NONE, +} + +PROC_STATUSES = { + cext.SIDL: _common.STATUS_IDLE, + cext.SRUN: _common.STATUS_RUNNING, + cext.SSLEEP: _common.STATUS_SLEEPING, + cext.SSTOP: _common.STATUS_STOPPED, + cext.SZOMB: _common.STATUS_ZOMBIE, +} + +kinfo_proc_map = dict( + ppid=0, + ruid=1, + euid=2, + suid=3, + rgid=4, + egid=5, + sgid=6, + ttynr=7, + ctime=8, + status=9, + name=10, +) + +pidtaskinfo_map = dict( + cpuutime=0, + cpustime=1, + rss=2, + vms=3, + pfaults=4, + pageins=5, + numthreads=6, + volctxsw=7, +) + + +# ===================================================================== +# --- memory +# ===================================================================== + + +def virtual_memory(): + """System virtual memory as a namedtuple.""" + total, active, inactive, wired, free, speculative = cext.virtual_mem() + # This is how Zabbix calculate avail and used mem: + # https://github.com/zabbix/zabbix/blob/master/src/libs/zbxsysinfo/osx/memory.c + # Also see: https://github.com/giampaolo/psutil/issues/1277 + avail = inactive + free + used = active + wired + # This is NOT how Zabbix calculates free mem but it matches "free" + # cmdline utility. + free -= speculative + percent = usage_percent((total - avail), total, round_=1) + return ntp.svmem( + total, avail, percent, used, free, active, inactive, wired + ) + + +def swap_memory(): + """Swap system memory as a (total, used, free, sin, sout) tuple.""" + total, used, free, sin, sout = cext.swap_mem() + percent = usage_percent(used, total, round_=1) + return ntp.sswap(total, used, free, percent, sin, sout) + + +# malloc / heap functions +heap_info = cext.heap_info +heap_trim = cext.heap_trim + + +# ===================================================================== +# --- CPU +# ===================================================================== + + +def cpu_times(): + """Return system CPU times as a namedtuple.""" + user, nice, system, idle = cext.cpu_times() + return ntp.scputimes(user, nice, system, idle) + + +def per_cpu_times(): + """Return system CPU times as a named tuple.""" + ret = [] + for cpu_t in cext.per_cpu_times(): + user, nice, system, idle = cpu_t + item = ntp.scputimes(user, nice, system, idle) + ret.append(item) + return ret + + +def cpu_count_logical(): + """Return the number of logical CPUs in the system.""" + return cext.cpu_count_logical() + + +def cpu_count_cores(): + """Return the number of CPU cores in the system.""" + return cext.cpu_count_cores() + + +def cpu_stats(): + ctx_switches, interrupts, soft_interrupts, syscalls, _traps = ( + cext.cpu_stats() + ) + return ntp.scpustats(ctx_switches, interrupts, soft_interrupts, syscalls) + + +if cext.has_cpu_freq(): # not always available on ARM64 + + def cpu_freq(): + """Return CPU frequency. + On macOS per-cpu frequency is not supported. + Also, the returned frequency never changes, see: + https://arstechnica.com/civis/viewtopic.php?f=19&t=465002. + """ + curr, min_, max_ = cext.cpu_freq() + return [ntp.scpufreq(curr, min_, max_)] + + +# ===================================================================== +# --- disks +# ===================================================================== + + +disk_usage = _psposix.disk_usage +disk_io_counters = cext.disk_io_counters + + +def disk_partitions(all=False): + """Return mounted disk partitions as a list of namedtuples.""" + retlist = [] + partitions = cext.disk_partitions() + for partition in partitions: + device, mountpoint, fstype, opts = partition + if device == 'none': + device = '' + if not all: + if not os.path.isabs(device) or not os.path.exists(device): + continue + ntuple = ntp.sdiskpart(device, mountpoint, fstype, opts) + retlist.append(ntuple) + return retlist + + +# ===================================================================== +# --- sensors +# ===================================================================== + + +def sensors_battery(): + """Return battery information.""" + try: + percent, minsleft, power_plugged = cext.sensors_battery() + except NotImplementedError: + # no power source - return None according to interface + return None + power_plugged = power_plugged == 1 + if power_plugged: + secsleft = _common.POWER_TIME_UNLIMITED + elif minsleft == -1: + secsleft = _common.POWER_TIME_UNKNOWN + else: + secsleft = minsleft * 60 + return ntp.sbattery(percent, secsleft, power_plugged) + + +# ===================================================================== +# --- network +# ===================================================================== + + +net_io_counters = cext.net_io_counters +net_if_addrs = cext.net_if_addrs + + +def net_connections(kind='inet'): + """System-wide network connections.""" + # Note: on macOS this will fail with AccessDenied unless + # the process is owned by root. + ret = [] + for pid in pids(): + try: + cons = Process(pid).net_connections(kind) + except NoSuchProcess: + continue + else: + if cons: + for c in cons: + c = list(c) + [pid] + ret.append(ntp.sconn(*c)) + return ret + + +def net_if_stats(): + """Get NIC stats (isup, duplex, speed, mtu).""" + names = net_io_counters().keys() + ret = {} + for name in names: + try: + mtu = cext.net_if_mtu(name) + flags = cext.net_if_flags(name) + duplex, speed = cext.net_if_duplex_speed(name) + except OSError as err: + # https://github.com/giampaolo/psutil/issues/1279 + if err.errno != errno.ENODEV: + raise + else: + if hasattr(_common, 'NicDuplex'): + duplex = _common.NicDuplex(duplex) + output_flags = ','.join(flags) + isup = 'running' in flags + ret[name] = ntp.snicstats(isup, duplex, speed, mtu, output_flags) + return ret + + +# ===================================================================== +# --- other system functions +# ===================================================================== + + +def boot_time(): + """The system boot time expressed in seconds since the epoch.""" + return cext.boot_time() + + +try: + INIT_BOOT_TIME = boot_time() +except Exception as err: # noqa: BLE001 + # Don't want to crash at import time. + debug(f"ignoring exception on import: {err!r}") + INIT_BOOT_TIME = 0 + + +def adjust_proc_create_time(ctime): + """Account for system clock updates.""" + if INIT_BOOT_TIME == 0: + return ctime + + diff = INIT_BOOT_TIME - boot_time() + if diff == 0 or abs(diff) < 1: + return ctime + + debug("system clock was updated; adjusting process create_time()") + if diff < 0: + return ctime - diff + return ctime + diff + + +def users(): + """Return currently connected users as a list of namedtuples.""" + retlist = [] + rawlist = cext.users() + for item in rawlist: + user, tty, hostname, tstamp, pid = item + if tty == '~': + continue # reboot or shutdown + if not tstamp: + continue + nt = ntp.suser(user, tty or None, hostname or None, tstamp, pid) + retlist.append(nt) + return retlist + + +# ===================================================================== +# --- processes +# ===================================================================== + + +def pids(): + ls = cext.pids() + if 0 not in ls: + # On certain macOS versions pids() C doesn't return PID 0 but + # "ps" does and the process is querable via sysctl(): + # https://travis-ci.org/giampaolo/psutil/jobs/309619941 + try: + Process(0).create_time() + ls.insert(0, 0) + except NoSuchProcess: + pass + except AccessDenied: + ls.insert(0, 0) + return ls + + +pid_exists = _psposix.pid_exists + + +def wrap_exceptions(fun): + """Decorator which translates bare OSError exceptions into + NoSuchProcess and AccessDenied. + """ + + @functools.wraps(fun) + def wrapper(self, *args, **kwargs): + pid, ppid, name = self.pid, self._ppid, self._name + try: + return fun(self, *args, **kwargs) + except ProcessLookupError as err: + if cext.proc_is_zombie(pid): + raise ZombieProcess(pid, name, ppid) from err + raise NoSuchProcess(pid, name) from err + except PermissionError as err: + raise AccessDenied(pid, name) from err + except cext.ZombieProcessError as err: + raise ZombieProcess(pid, name, ppid) from err + + return wrapper + + +class Process: + """Wrapper class around underlying C implementation.""" + + __slots__ = ["_cache", "_name", "_ppid", "pid"] + + def __init__(self, pid): + self.pid = pid + self._name = None + self._ppid = None + + @wrap_exceptions + @memoize_when_activated + def _get_kinfo_proc(self): + # Note: should work with all PIDs without permission issues. + ret = cext.proc_kinfo_oneshot(self.pid) + assert len(ret) == len(kinfo_proc_map) + return ret + + @wrap_exceptions + @memoize_when_activated + def _get_pidtaskinfo(self): + # Note: should work for PIDs owned by user only. + ret = cext.proc_pidtaskinfo_oneshot(self.pid) + assert len(ret) == len(pidtaskinfo_map) + return ret + + def oneshot_enter(self): + self._get_kinfo_proc.cache_activate(self) + self._get_pidtaskinfo.cache_activate(self) + + def oneshot_exit(self): + self._get_kinfo_proc.cache_deactivate(self) + self._get_pidtaskinfo.cache_deactivate(self) + + @wrap_exceptions + def name(self): + name = self._get_kinfo_proc()[kinfo_proc_map['name']] + return name if name is not None else cext.proc_name(self.pid) + + @wrap_exceptions + def exe(self): + return cext.proc_exe(self.pid) + + @wrap_exceptions + def cmdline(self): + return cext.proc_cmdline(self.pid) + + @wrap_exceptions + def environ(self): + return parse_environ_block(cext.proc_environ(self.pid)) + + @wrap_exceptions + def ppid(self): + self._ppid = self._get_kinfo_proc()[kinfo_proc_map['ppid']] + return self._ppid + + @wrap_exceptions + def cwd(self): + return cext.proc_cwd(self.pid) + + @wrap_exceptions + def uids(self): + rawtuple = self._get_kinfo_proc() + return ntp.puids( + rawtuple[kinfo_proc_map['ruid']], + rawtuple[kinfo_proc_map['euid']], + rawtuple[kinfo_proc_map['suid']], + ) + + @wrap_exceptions + def gids(self): + rawtuple = self._get_kinfo_proc() + return ntp.puids( + rawtuple[kinfo_proc_map['rgid']], + rawtuple[kinfo_proc_map['egid']], + rawtuple[kinfo_proc_map['sgid']], + ) + + @wrap_exceptions + def terminal(self): + tty_nr = self._get_kinfo_proc()[kinfo_proc_map['ttynr']] + tmap = _psposix.get_terminal_map() + try: + return tmap[tty_nr] + except KeyError: + return None + + @wrap_exceptions + def memory_info(self): + rawtuple = self._get_pidtaskinfo() + return ntp.pmem( + rawtuple[pidtaskinfo_map['rss']], + rawtuple[pidtaskinfo_map['vms']], + rawtuple[pidtaskinfo_map['pfaults']], + rawtuple[pidtaskinfo_map['pageins']], + ) + + @wrap_exceptions + def memory_full_info(self): + basic_mem = self.memory_info() + uss = cext.proc_memory_uss(self.pid) + return ntp.pfullmem(*basic_mem + (uss,)) + + @wrap_exceptions + def cpu_times(self): + rawtuple = self._get_pidtaskinfo() + return ntp.pcputimes( + rawtuple[pidtaskinfo_map['cpuutime']], + rawtuple[pidtaskinfo_map['cpustime']], + # children user / system times are not retrievable (set to 0) + 0.0, + 0.0, + ) + + @wrap_exceptions + def create_time(self, monotonic=False): + ctime = self._get_kinfo_proc()[kinfo_proc_map['ctime']] + if not monotonic: + ctime = adjust_proc_create_time(ctime) + return ctime + + @wrap_exceptions + def num_ctx_switches(self): + # Unvoluntary value seems not to be available; + # getrusage() numbers seems to confirm this theory. + # We set it to 0. + vol = self._get_pidtaskinfo()[pidtaskinfo_map['volctxsw']] + return ntp.pctxsw(vol, 0) + + @wrap_exceptions + def num_threads(self): + return self._get_pidtaskinfo()[pidtaskinfo_map['numthreads']] + + @wrap_exceptions + def open_files(self): + if self.pid == 0: + return [] + files = [] + rawlist = cext.proc_open_files(self.pid) + for path, fd in rawlist: + if isfile_strict(path): + ntuple = ntp.popenfile(path, fd) + files.append(ntuple) + return files + + @wrap_exceptions + def net_connections(self, kind='inet'): + families, types = conn_tmap[kind] + rawlist = cext.proc_net_connections(self.pid, families, types) + ret = [] + for item in rawlist: + fd, fam, type, laddr, raddr, status = item + nt = conn_to_ntuple( + fd, fam, type, laddr, raddr, status, TCP_STATUSES + ) + ret.append(nt) + return ret + + @wrap_exceptions + def num_fds(self): + if self.pid == 0: + return 0 + return cext.proc_num_fds(self.pid) + + @wrap_exceptions + def wait(self, timeout=None): + return _psposix.wait_pid(self.pid, timeout) + + @wrap_exceptions + def nice_get(self): + return cext.proc_priority_get(self.pid) + + @wrap_exceptions + def nice_set(self, value): + return cext.proc_priority_set(self.pid, value) + + @wrap_exceptions + def status(self): + code = self._get_kinfo_proc()[kinfo_proc_map['status']] + # XXX is '?' legit? (we're not supposed to return it anyway) + return PROC_STATUSES.get(code, '?') + + @wrap_exceptions + def threads(self): + rawlist = cext.proc_threads(self.pid) + retlist = [] + for thread_id, utime, stime in rawlist: + ntuple = ntp.pthread(thread_id, utime, stime) + retlist.append(ntuple) + return retlist diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/psutil/_psposix.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/psutil/_psposix.py new file mode 100644 index 0000000000000000000000000000000000000000..c7ad2fdabd30ee585642b49e38db9afc489b6ede --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/psutil/_psposix.py @@ -0,0 +1,363 @@ +# Copyright (c) 2009, Giampaolo Rodola'. All rights reserved. +# Use of this source code is governed by a BSD-style license that can be +# found in the LICENSE file. + +"""Routines common to all posix systems.""" + +import enum +import errno +import glob +import os +import select +import signal +import time + +from . import _ntuples as ntp +from ._common import MACOS +from ._common import TimeoutExpired +from ._common import debug +from ._common import memoize +from ._common import usage_percent + +if MACOS: + from . import _psutil_osx + + +__all__ = ['pid_exists', 'wait_pid', 'disk_usage', 'get_terminal_map'] + + +def pid_exists(pid): + """Check whether pid exists in the current process table.""" + if pid == 0: + # According to "man 2 kill" PID 0 has a special meaning: + # it refers to <> so we don't want to go any further. + # If we get here it means this UNIX platform *does* have + # a process with id 0. + return True + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + except PermissionError: + # EPERM clearly means there's a process to deny access to + return True + # According to "man 2 kill" possible error values are + # (EINVAL, EPERM, ESRCH) + else: + return True + + +Negsignal = enum.IntEnum( + 'Negsignal', {x.name: -x.value for x in signal.Signals} +) + + +def negsig_to_enum(num): + """Convert a negative signal value to an enum.""" + try: + return Negsignal(num) + except ValueError: + return num + + +def convert_exit_code(status): + """Convert a os.waitpid() status to an exit code.""" + if os.WIFEXITED(status): + # Process terminated normally by calling exit(3) or _exit(2), + # or by returning from main(). The return value is the + # positive integer passed to *exit(). + return os.WEXITSTATUS(status) + if os.WIFSIGNALED(status): + # Process exited due to a signal. Return the negative value + # of that signal. + return negsig_to_enum(-os.WTERMSIG(status)) + # if os.WIFSTOPPED(status): + # # Process was stopped via SIGSTOP or is being traced, and + # # waitpid() was called with WUNTRACED flag. PID is still + # # alive. From now on waitpid() will keep returning (0, 0) + # # until the process state doesn't change. + # # It may make sense to catch/enable this since stopped PIDs + # # ignore SIGTERM. + # interval = sleep(interval) + # continue + # if os.WIFCONTINUED(status): + # # Process was resumed via SIGCONT and waitpid() was called + # # with WCONTINUED flag. + # interval = sleep(interval) + # continue + + # Should never happen. + msg = f"unknown process exit status {status!r}" + raise ValueError(msg) + + +def wait_pid_posix( + pid, + timeout=None, + _waitpid=os.waitpid, + _timer=getattr(time, 'monotonic', time.time), # noqa: B008 + _min=min, + _sleep=time.sleep, + _pid_exists=pid_exists, +): + """Wait for a process PID to terminate. + + If the process terminated normally by calling exit(3) or _exit(2), + or by returning from main(), the return value is the positive integer + passed to *exit(). + + If it was terminated by a signal it returns the negated value of the + signal which caused the termination (e.g. -SIGTERM). + + If PID is not a children of os.getpid() (current process) just + wait until the process disappears and return None. + + If PID does not exist at all return None immediately. + + If timeout is specified and process is still alive raise + TimeoutExpired. + + If timeout=0 either return immediately or raise TimeoutExpired + (non-blocking). + """ + interval = 0.0001 + max_interval = 0.04 + flags = 0 + stop_at = None + + if timeout is not None: + flags |= os.WNOHANG + if timeout != 0: + stop_at = _timer() + timeout + + def sleep_or_timeout(interval): + # Sleep for some time and return a new increased interval. + if timeout == 0 or (stop_at is not None and _timer() >= stop_at): + raise TimeoutExpired(timeout) + _sleep(interval) + return _min(interval * 2, max_interval) + + # See: https://linux.die.net/man/2/waitpid + while True: + try: + retpid, status = os.waitpid(pid, flags) + except ChildProcessError: + # This has two meanings: + # - PID is not a child of os.getpid() in which case + # we keep polling until it's gone + # - PID never existed in the first place + # In both cases we'll eventually return None as we + # can't determine its exit status code. + while _pid_exists(pid): + interval = sleep_or_timeout(interval) + return None + else: + if retpid == 0: + # WNOHANG flag was used and PID is still running. + interval = sleep_or_timeout(interval) + else: + return convert_exit_code(status) + + +def _waitpid(pid, timeout): + """Wrapper around os.waitpid(). PID is supposed to be gone already, + it just returns the exit code. + """ + try: + retpid, status = os.waitpid(pid, 0) + except ChildProcessError: + # PID is not a child of os.getpid(). + return wait_pid_posix(pid, timeout) + else: + assert retpid != 0 + return convert_exit_code(status) + + +def wait_pid_pidfd_open(pid, timeout=None): + """Wait for PID to terminate using pidfd_open() + poll(). Linux >= + 5.3 + Python >= 3.9 only. + """ + try: + pidfd = os.pidfd_open(pid, 0) + except OSError as err: + if err.errno == errno.ESRCH: + # No such process. os.waitpid() may still be able to return + # the status code. + return wait_pid_posix(pid, timeout) + if err.errno in {errno.EMFILE, errno.ENFILE, errno.ENODEV}: + # EMFILE, ENFILE: too many open files + # ENODEV: anonymous inode filesystem not supported + debug(f"pidfd_open() failed ({err!r}); use fallback") + return wait_pid_posix(pid, timeout) + raise + + try: + # poll() / select() have the advantage of not requiring any + # extra file descriptor, contrary to epoll() / kqueue(). + # select() crashes if process opens > 1024 FDs, so we use + # poll(). + poller = select.poll() + poller.register(pidfd, select.POLLIN) + timeout_ms = None if timeout is None else int(timeout * 1000) + events = poller.poll(timeout_ms) # wait + + if not events: + raise TimeoutExpired(timeout) + return _waitpid(pid, timeout) + finally: + os.close(pidfd) + + +def wait_pid_kqueue(pid, timeout=None): + """Wait for PID to terminate using kqueue(). macOS and BSD only.""" + try: + kq = select.kqueue() + except OSError as err: + if err.errno in {errno.EMFILE, errno.ENFILE}: # too many open files + debug(f"kqueue() failed ({err!r}); use fallback") + return wait_pid_posix(pid, timeout) + raise + + try: + kev = select.kevent( + pid, + filter=select.KQ_FILTER_PROC, + flags=select.KQ_EV_ADD | select.KQ_EV_ONESHOT, + fflags=select.KQ_NOTE_EXIT, + ) + try: + events = kq.control([kev], 1, timeout) # wait + except OSError as err: + if err.errno in {errno.EACCES, errno.EPERM, errno.ESRCH}: + debug(f"kqueue.control() failed ({err!r}); use fallback") + return wait_pid_posix(pid, timeout) + raise + else: + if not events: + raise TimeoutExpired(timeout) + return _waitpid(pid, timeout) + finally: + kq.close() + + +@memoize +def can_use_pidfd_open(): + # Availability: Linux >= 5.3, Python >= 3.9 + if not hasattr(os, "pidfd_open"): + return False + try: + pidfd = os.pidfd_open(os.getpid(), 0) + except OSError as err: + if err.errno in {errno.EMFILE, errno.ENFILE}: # noqa: SIM103 + # transitory 'too many open files' + return True + # likely blocked by security policy like SECCOMP (EPERM, + # EACCES, ENOSYS) + return False + else: + os.close(pidfd) + return True + + +@memoize +def can_use_kqueue(): + # Availability: macOS, BSD + names = ( + "kqueue", + "KQ_EV_ADD", + "KQ_EV_ONESHOT", + "KQ_FILTER_PROC", + "KQ_NOTE_EXIT", + ) + if not all(hasattr(select, x) for x in names): + return False + kq = None + try: + kq = select.kqueue() + kev = select.kevent( + os.getpid(), + filter=select.KQ_FILTER_PROC, + flags=select.KQ_EV_ADD | select.KQ_EV_ONESHOT, + fflags=select.KQ_NOTE_EXIT, + ) + kq.control([kev], 1, 0) + return True + except OSError as err: + if err.errno in {errno.EMFILE, errno.ENFILE}: # noqa: SIM103 + # transitory 'too many open files' + return True + return False + finally: + if kq is not None: + kq.close() + + +def wait_pid(pid, timeout=None): + # PID 0 passed to waitpid() waits for any child of the current + # process to change state. + assert pid > 0 + if timeout is not None: + assert timeout >= 0 + + if can_use_pidfd_open(): + return wait_pid_pidfd_open(pid, timeout) + elif can_use_kqueue(): + return wait_pid_kqueue(pid, timeout) + else: + return wait_pid_posix(pid, timeout) + + +wait_pid.__doc__ = wait_pid_posix.__doc__ + + +def disk_usage(path): + """Return disk usage associated with path. + Note: UNIX usually reserves 5% disk space which is not accessible + by user. In this function "total" and "used" values reflect the + total and used disk space whereas "free" and "percent" represent + the "free" and "used percent" user disk space. + """ + st = os.statvfs(path) + # Total space which is only available to root (unless changed + # at system level). + total = st.f_blocks * st.f_frsize + # Remaining free space usable by root. + avail_to_root = st.f_bfree * st.f_frsize + # Remaining free space usable by user. + avail_to_user = st.f_bavail * st.f_frsize + # Total space being used in general. + used = total - avail_to_root + if MACOS: + # see: https://github.com/giampaolo/psutil/pull/2152 + used = _psutil_osx.disk_usage_used(path, used) + # Total space which is available to user (same as 'total' but + # for the user). + total_user = used + avail_to_user + # User usage percent compared to the total amount of space + # the user can use. This number would be higher if compared + # to root's because the user has less space (usually -5%). + usage_percent_user = usage_percent(used, total_user, round_=1) + + # NB: the percentage is -5% than what shown by df due to + # reserved blocks that we are currently not considering: + # https://github.com/giampaolo/psutil/issues/829#issuecomment-223750462 + return ntp.sdiskusage( + total=total, used=used, free=avail_to_user, percent=usage_percent_user + ) + + +@memoize +def get_terminal_map(): + """Get a map of device-id -> path as a dict. + Used by Process.terminal(). + """ + ret = {} + ls = glob.glob('/dev/tty*') + glob.glob('/dev/pts/*') + for name in ls: + assert name not in ret, name + try: + ret[os.stat(name).st_rdev] = name + except FileNotFoundError: + pass + return ret diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/psutil/_pssunos.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/psutil/_pssunos.py new file mode 100644 index 0000000000000000000000000000000000000000..a0e70581f5a02dfbc8e61a9d26e105ee6160a8aa --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/psutil/_pssunos.py @@ -0,0 +1,704 @@ +# Copyright (c) 2009, Giampaolo Rodola'. All rights reserved. +# Use of this source code is governed by a BSD-style license that can be +# found in the LICENSE file. + +"""Sun OS Solaris platform implementation.""" + +import errno +import functools +import os +import socket +import subprocess +import sys +from collections import namedtuple +from socket import AF_INET + +from . import _common +from . import _ntuples as ntp +from . import _psposix +from . import _psutil_sunos as cext +from ._common import AF_INET6 +from ._common import ENCODING +from ._common import AccessDenied +from ._common import NoSuchProcess +from ._common import ZombieProcess +from ._common import debug +from ._common import get_procfs_path +from ._common import isfile_strict +from ._common import memoize_when_activated +from ._common import sockfam_to_enum +from ._common import socktype_to_enum +from ._common import usage_percent + +__extra__all__ = ["CONN_IDLE", "CONN_BOUND", "PROCFS_PATH"] + + +# ===================================================================== +# --- globals +# ===================================================================== + + +PAGE_SIZE = cext.getpagesize() +AF_LINK = cext.AF_LINK +IS_64_BIT = sys.maxsize > 2**32 + +CONN_IDLE = "IDLE" +CONN_BOUND = "BOUND" + +PROC_STATUSES = { + cext.SSLEEP: _common.STATUS_SLEEPING, + cext.SRUN: _common.STATUS_RUNNING, + cext.SZOMB: _common.STATUS_ZOMBIE, + cext.SSTOP: _common.STATUS_STOPPED, + cext.SIDL: _common.STATUS_IDLE, + cext.SONPROC: _common.STATUS_RUNNING, # same as run + cext.SWAIT: _common.STATUS_WAITING, +} + +TCP_STATUSES = { + cext.TCPS_ESTABLISHED: _common.CONN_ESTABLISHED, + cext.TCPS_SYN_SENT: _common.CONN_SYN_SENT, + cext.TCPS_SYN_RCVD: _common.CONN_SYN_RECV, + cext.TCPS_FIN_WAIT_1: _common.CONN_FIN_WAIT1, + cext.TCPS_FIN_WAIT_2: _common.CONN_FIN_WAIT2, + cext.TCPS_TIME_WAIT: _common.CONN_TIME_WAIT, + cext.TCPS_CLOSED: _common.CONN_CLOSE, + cext.TCPS_CLOSE_WAIT: _common.CONN_CLOSE_WAIT, + cext.TCPS_LAST_ACK: _common.CONN_LAST_ACK, + cext.TCPS_LISTEN: _common.CONN_LISTEN, + cext.TCPS_CLOSING: _common.CONN_CLOSING, + cext.PSUTIL_CONN_NONE: _common.CONN_NONE, + cext.TCPS_IDLE: CONN_IDLE, # sunos specific + cext.TCPS_BOUND: CONN_BOUND, # sunos specific +} + +proc_info_map = dict( + ppid=0, + rss=1, + vms=2, + create_time=3, + nice=4, + num_threads=5, + status=6, + ttynr=7, + uid=8, + euid=9, + gid=10, + egid=11, +) + + +# ===================================================================== +# --- memory +# ===================================================================== + + +def virtual_memory(): + """Report virtual memory metrics.""" + # we could have done this with kstat, but IMHO this is good enough + total = os.sysconf('SC_PHYS_PAGES') * PAGE_SIZE + # note: there's no difference on Solaris + free = avail = os.sysconf('SC_AVPHYS_PAGES') * PAGE_SIZE + used = total - free + percent = usage_percent(used, total, round_=1) + return ntp.svmem(total, avail, percent, used, free) + + +def swap_memory(): + """Report swap memory metrics.""" + sin, sout = cext.swap_mem() + # XXX + # we are supposed to get total/free by doing so: + # http://cvs.opensolaris.org/source/xref/onnv/onnv-gate/ + # usr/src/cmd/swap/swap.c + # ...nevertheless I can't manage to obtain the same numbers as 'swap' + # cmdline utility, so let's parse its output (sigh!) + p = subprocess.Popen( + [ + '/usr/bin/env', + f"PATH=/usr/sbin:/sbin:{os.environ['PATH']}", + 'swap', + '-l', + ], + stdout=subprocess.PIPE, + ) + stdout, _ = p.communicate() + stdout = stdout.decode(sys.stdout.encoding) + if p.returncode != 0: + msg = f"'swap -l' failed (retcode={p.returncode})" + raise RuntimeError(msg) + + lines = stdout.strip().split('\n')[1:] + if not lines: + msg = 'no swap device(s) configured' + raise RuntimeError(msg) + total = free = 0 + for line in lines: + line = line.split() + t, f = line[3:5] + total += int(int(t) * 512) + free += int(int(f) * 512) + used = total - free + percent = usage_percent(used, total, round_=1) + return ntp.sswap( + total, used, free, percent, sin * PAGE_SIZE, sout * PAGE_SIZE + ) + + +# ===================================================================== +# --- CPU +# ===================================================================== + + +def cpu_times(): + """Return system-wide CPU times as a named tuple.""" + ret = cext.per_cpu_times() + return ntp.scputimes(*[sum(x) for x in zip(*ret)]) + + +def per_cpu_times(): + """Return system per-CPU times as a list of named tuples.""" + ret = cext.per_cpu_times() + return [ntp.scputimes(*x) for x in ret] + + +def cpu_count_logical(): + """Return the number of logical CPUs in the system.""" + try: + return os.sysconf("SC_NPROCESSORS_ONLN") + except ValueError: + # mimic os.cpu_count() behavior + return None + + +def cpu_count_cores(): + """Return the number of CPU cores in the system.""" + return cext.cpu_count_cores() + + +def cpu_stats(): + """Return various CPU stats as a named tuple.""" + ctx_switches, interrupts, syscalls, _traps = cext.cpu_stats() + soft_interrupts = 0 + return ntp.scpustats(ctx_switches, interrupts, soft_interrupts, syscalls) + + +# ===================================================================== +# --- disks +# ===================================================================== + + +disk_io_counters = cext.disk_io_counters +disk_usage = _psposix.disk_usage + + +def disk_partitions(all=False): + """Return system disk partitions.""" + # TODO - the filtering logic should be better checked so that + # it tries to reflect 'df' as much as possible + retlist = [] + partitions = cext.disk_partitions() + for partition in partitions: + device, mountpoint, fstype, opts = partition + if device == 'none': + device = '' + if not all: + # Differently from, say, Linux, we don't have a list of + # common fs types so the best we can do, AFAIK, is to + # filter by filesystem having a total size > 0. + try: + if not disk_usage(mountpoint).total: + continue + except OSError as err: + # https://github.com/giampaolo/psutil/issues/1674 + debug(f"skipping {mountpoint!r}: {err}") + continue + ntuple = ntp.sdiskpart(device, mountpoint, fstype, opts) + retlist.append(ntuple) + return retlist + + +# ===================================================================== +# --- network +# ===================================================================== + + +net_io_counters = cext.net_io_counters +net_if_addrs = cext.net_if_addrs + + +def net_connections(kind, _pid=-1): + """Return socket connections. If pid == -1 return system-wide + connections (as opposed to connections opened by one process only). + Only INET sockets are returned (UNIX are not). + """ + families, types = _common.conn_tmap[kind] + rawlist = cext.net_connections(_pid) + ret = set() + for item in rawlist: + fd, fam, type_, laddr, raddr, status, pid = item + if fam not in families: + continue + if type_ not in types: + continue + # TODO: refactor and use _common.conn_to_ntuple. + if fam in {AF_INET, AF_INET6}: + if laddr: + laddr = ntp.addr(*laddr) + if raddr: + raddr = ntp.addr(*raddr) + status = TCP_STATUSES[status] + fam = sockfam_to_enum(fam) + type_ = socktype_to_enum(type_) + if _pid == -1: + nt = ntp.sconn(fd, fam, type_, laddr, raddr, status, pid) + else: + nt = ntp.pconn(fd, fam, type_, laddr, raddr, status) + ret.add(nt) + return list(ret) + + +def net_if_stats(): + """Get NIC stats (isup, duplex, speed, mtu).""" + ret = cext.net_if_stats() + for name, items in ret.items(): + isup, duplex, speed, mtu = items + if hasattr(_common, 'NicDuplex'): + duplex = _common.NicDuplex(duplex) + ret[name] = ntp.snicstats(isup, duplex, speed, mtu, '') + return ret + + +# ===================================================================== +# --- other system functions +# ===================================================================== + + +def boot_time(): + """The system boot time expressed in seconds since the epoch.""" + return cext.boot_time() + + +def users(): + """Return currently connected users as a list of namedtuples.""" + retlist = [] + rawlist = cext.users() + localhost = (':0.0', ':0') + for item in rawlist: + user, tty, hostname, tstamp, user_process, pid = item + # note: the underlying C function includes entries about + # system boot, run level and others. We might want + # to use them in the future. + if not user_process: + continue + if hostname in localhost: + hostname = 'localhost' + nt = ntp.suser(user, tty, hostname, tstamp, pid) + retlist.append(nt) + return retlist + + +# ===================================================================== +# --- processes +# ===================================================================== + + +def pids(): + """Returns a list of PIDs currently running on the system.""" + path = get_procfs_path().encode(ENCODING) + return [int(x) for x in os.listdir(path) if x.isdigit()] + + +def pid_exists(pid): + """Check for the existence of a unix pid.""" + return _psposix.pid_exists(pid) + + +def wrap_exceptions(fun): + """Call callable into a try/except clause and translate ENOENT, + EACCES and EPERM in NoSuchProcess or AccessDenied exceptions. + """ + + @functools.wraps(fun) + def wrapper(self, *args, **kwargs): + pid, ppid, name = self.pid, self._ppid, self._name + try: + return fun(self, *args, **kwargs) + except (FileNotFoundError, ProcessLookupError) as err: + # ENOENT (no such file or directory) gets raised on open(). + # ESRCH (no such process) can get raised on read() if + # process is gone in meantime. + if not pid_exists(pid): + raise NoSuchProcess(pid, name) from err + raise ZombieProcess(pid, name, ppid) from err + except PermissionError as err: + raise AccessDenied(pid, name) from err + except OSError as err: + if pid == 0: + if 0 in pids(): + raise AccessDenied(pid, name) from err + raise + raise + + return wrapper + + +class Process: + """Wrapper class around underlying C implementation.""" + + __slots__ = ["_cache", "_name", "_ppid", "_procfs_path", "pid"] + + def __init__(self, pid): + self.pid = pid + self._name = None + self._ppid = None + self._procfs_path = get_procfs_path() + + def _assert_alive(self): + """Raise NSP if the process disappeared on us.""" + # For those C function who do not raise NSP, possibly returning + # incorrect or incomplete result. + os.stat(f"{self._procfs_path}/{self.pid}") + + def oneshot_enter(self): + self._proc_name_and_args.cache_activate(self) + self._proc_basic_info.cache_activate(self) + self._proc_cred.cache_activate(self) + + def oneshot_exit(self): + self._proc_name_and_args.cache_deactivate(self) + self._proc_basic_info.cache_deactivate(self) + self._proc_cred.cache_deactivate(self) + + @wrap_exceptions + @memoize_when_activated + def _proc_name_and_args(self): + return cext.proc_name_and_args(self.pid, self._procfs_path) + + @wrap_exceptions + @memoize_when_activated + def _proc_basic_info(self): + if self.pid == 0 and not os.path.exists( + f"{self._procfs_path}/{self.pid}/psinfo" + ): + raise AccessDenied(self.pid) + ret = cext.proc_basic_info(self.pid, self._procfs_path) + assert len(ret) == len(proc_info_map) + return ret + + @wrap_exceptions + @memoize_when_activated + def _proc_cred(self): + return cext.proc_cred(self.pid, self._procfs_path) + + @wrap_exceptions + def name(self): + # note: max len == 15 + return self._proc_name_and_args()[0] + + @wrap_exceptions + def exe(self): + try: + return os.readlink(f"{self._procfs_path}/{self.pid}/path/a.out") + except OSError: + pass # continue and guess the exe name from the cmdline + # Will be guessed later from cmdline but we want to explicitly + # invoke cmdline here in order to get an AccessDenied + # exception if the user has not enough privileges. + self.cmdline() + return "" + + @wrap_exceptions + def cmdline(self): + return self._proc_name_and_args()[1] + + @wrap_exceptions + def environ(self): + return cext.proc_environ(self.pid, self._procfs_path) + + @wrap_exceptions + def create_time(self): + return self._proc_basic_info()[proc_info_map['create_time']] + + @wrap_exceptions + def num_threads(self): + return self._proc_basic_info()[proc_info_map['num_threads']] + + @wrap_exceptions + def nice_get(self): + # Note #1: getpriority(3) doesn't work for realtime processes. + # Psinfo is what ps uses, see: + # https://github.com/giampaolo/psutil/issues/1194 + return self._proc_basic_info()[proc_info_map['nice']] + + @wrap_exceptions + def nice_set(self, value): + if self.pid in {2, 3}: + # Special case PIDs: internally setpriority(3) return ESRCH + # (no such process), no matter what. + # The process actually exists though, as it has a name, + # creation time, etc. + raise AccessDenied(self.pid, self._name) + return cext.proc_priority_set(self.pid, value) + + @wrap_exceptions + def ppid(self): + self._ppid = self._proc_basic_info()[proc_info_map['ppid']] + return self._ppid + + @wrap_exceptions + def uids(self): + try: + real, effective, saved, _, _, _ = self._proc_cred() + except AccessDenied: + real = self._proc_basic_info()[proc_info_map['uid']] + effective = self._proc_basic_info()[proc_info_map['euid']] + saved = None + return ntp.puids(real, effective, saved) + + @wrap_exceptions + def gids(self): + try: + _, _, _, real, effective, saved = self._proc_cred() + except AccessDenied: + real = self._proc_basic_info()[proc_info_map['gid']] + effective = self._proc_basic_info()[proc_info_map['egid']] + saved = None + return ntp.puids(real, effective, saved) + + @wrap_exceptions + def cpu_times(self): + try: + times = cext.proc_cpu_times(self.pid, self._procfs_path) + except OSError as err: + if err.errno == errno.EOVERFLOW and not IS_64_BIT: + # We may get here if we attempt to query a 64bit process + # with a 32bit python. + # Error originates from read() and also tools like "cat" + # fail in the same way (!). + # Since there simply is no way to determine CPU times we + # return 0.0 as a fallback. See: + # https://github.com/giampaolo/psutil/issues/857 + times = (0.0, 0.0, 0.0, 0.0) + else: + raise + return ntp.pcputimes(*times) + + @wrap_exceptions + def cpu_num(self): + return cext.proc_cpu_num(self.pid, self._procfs_path) + + @wrap_exceptions + def terminal(self): + procfs_path = self._procfs_path + hit_enoent = False + tty = wrap_exceptions(self._proc_basic_info()[proc_info_map['ttynr']]) + if tty != cext.PRNODEV: + for x in (0, 1, 2, 255): + try: + return os.readlink(f"{procfs_path}/{self.pid}/path/{x}") + except FileNotFoundError: + hit_enoent = True + continue + if hit_enoent: + self._assert_alive() + + @wrap_exceptions + def cwd(self): + # /proc/PID/path/cwd may not be resolved by readlink() even if + # it exists (ls shows it). If that's the case and the process + # is still alive return None (we can return None also on BSD). + # Reference: https://groups.google.com/g/comp.unix.solaris/c/tcqvhTNFCAs + procfs_path = self._procfs_path + try: + return os.readlink(f"{procfs_path}/{self.pid}/path/cwd") + except FileNotFoundError: + os.stat(f"{procfs_path}/{self.pid}") # raise NSP or AD + return "" + + @wrap_exceptions + def memory_info(self): + ret = self._proc_basic_info() + rss = ret[proc_info_map['rss']] * 1024 + vms = ret[proc_info_map['vms']] * 1024 + return ntp.pmem(rss, vms) + + memory_full_info = memory_info + + @wrap_exceptions + def status(self): + code = self._proc_basic_info()[proc_info_map['status']] + # XXX is '?' legit? (we're not supposed to return it anyway) + return PROC_STATUSES.get(code, '?') + + @wrap_exceptions + def threads(self): + procfs_path = self._procfs_path + ret = [] + tids = os.listdir(f"{procfs_path}/{self.pid}/lwp") + hit_enoent = False + for tid in tids: + tid = int(tid) + try: + utime, stime = cext.query_process_thread( + self.pid, tid, procfs_path + ) + except OSError as err: + if err.errno == errno.EOVERFLOW and not IS_64_BIT: + # We may get here if we attempt to query a 64bit process + # with a 32bit python. + # Error originates from read() and also tools like "cat" + # fail in the same way (!). + # Since there simply is no way to determine CPU times we + # return 0.0 as a fallback. See: + # https://github.com/giampaolo/psutil/issues/857 + continue + # ENOENT == thread gone in meantime + if err.errno == errno.ENOENT: + hit_enoent = True + continue + raise + else: + nt = ntp.pthread(tid, utime, stime) + ret.append(nt) + if hit_enoent: + self._assert_alive() + return ret + + @wrap_exceptions + def open_files(self): + retlist = [] + hit_enoent = False + procfs_path = self._procfs_path + pathdir = f"{procfs_path}/{self.pid}/path" + for fd in os.listdir(f"{procfs_path}/{self.pid}/fd"): + path = os.path.join(pathdir, fd) + if os.path.islink(path): + try: + file = os.readlink(path) + except FileNotFoundError: + hit_enoent = True + continue + else: + if isfile_strict(file): + retlist.append(ntp.popenfile(file, int(fd))) + if hit_enoent: + self._assert_alive() + return retlist + + def _get_unix_sockets(self, pid): + """Get UNIX sockets used by process by parsing 'pfiles' output.""" + # TODO: rewrite this in C (...but the damn netstat source code + # does not include this part! Argh!!) + cmd = ["pfiles", str(pid)] + p = subprocess.Popen( + cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE + ) + stdout, stderr = p.communicate() + stdout, stderr = ( + x.decode(sys.stdout.encoding) for x in (stdout, stderr) + ) + if p.returncode != 0: + if 'permission denied' in stderr.lower(): + raise AccessDenied(self.pid, self._name) + if 'no such process' in stderr.lower(): + raise NoSuchProcess(self.pid, self._name) + msg = f"{cmd!r} command error\n{stderr}" + raise RuntimeError(msg) + + lines = stdout.split('\n')[2:] + for i, line in enumerate(lines): + line = line.lstrip() + if line.startswith('sockname: AF_UNIX'): + path = line.split(' ', 2)[2] + type = lines[i - 2].strip() + if type == 'SOCK_STREAM': + type = socket.SOCK_STREAM + elif type == 'SOCK_DGRAM': + type = socket.SOCK_DGRAM + else: + type = -1 + yield (-1, socket.AF_UNIX, type, path, "", _common.CONN_NONE) + + @wrap_exceptions + def net_connections(self, kind='inet'): + ret = net_connections(kind, _pid=self.pid) + # The underlying C implementation retrieves all OS connections + # and filters them by PID. At this point we can't tell whether + # an empty list means there were no connections for process or + # process is no longer active so we force NSP in case the PID + # is no longer there. + if not ret: + # will raise NSP if process is gone + os.stat(f"{self._procfs_path}/{self.pid}") + + # UNIX sockets + if kind in {'all', 'unix'}: + ret.extend( + [ntp.pconn(*conn) for conn in self._get_unix_sockets(self.pid)] + ) + return ret + + nt_mmap_grouped = namedtuple('mmap', 'path rss anon locked') + nt_mmap_ext = namedtuple('mmap', 'addr perms path rss anon locked') + + @wrap_exceptions + def memory_maps(self): + def toaddr(start, end): + return "{}-{}".format( + hex(start)[2:].strip('L'), hex(end)[2:].strip('L') + ) + + procfs_path = self._procfs_path + retlist = [] + try: + rawlist = cext.proc_memory_maps(self.pid, procfs_path) + except OSError as err: + if err.errno == errno.EOVERFLOW and not IS_64_BIT: + # We may get here if we attempt to query a 64bit process + # with a 32bit python. + # Error originates from read() and also tools like "cat" + # fail in the same way (!). + # Since there simply is no way to determine CPU times we + # return 0.0 as a fallback. See: + # https://github.com/giampaolo/psutil/issues/857 + return [] + else: + raise + hit_enoent = False + for item in rawlist: + addr, addrsize, perm, name, rss, anon, locked = item + addr = toaddr(addr, addrsize) + if not name.startswith('['): + try: + name = os.readlink(f"{procfs_path}/{self.pid}/path/{name}") + except OSError as err: + if err.errno == errno.ENOENT: + # sometimes the link may not be resolved by + # readlink() even if it exists (ls shows it). + # If that's the case we just return the + # unresolved link path. + # This seems an inconsistency with /proc similar + # to: http://goo.gl/55XgO + name = f"{procfs_path}/{self.pid}/path/{name}" + hit_enoent = True + else: + raise + retlist.append((addr, perm, name, rss, anon, locked)) + if hit_enoent: + self._assert_alive() + return retlist + + @wrap_exceptions + def num_fds(self): + return len(os.listdir(f"{self._procfs_path}/{self.pid}/fd")) + + @wrap_exceptions + def num_ctx_switches(self): + return ntp.pctxsw( + *cext.proc_num_ctx_switches(self.pid, self._procfs_path) + ) + + @wrap_exceptions + def wait(self, timeout=None): + return _psposix.wait_pid(self.pid, timeout) diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/psutil/_pswindows.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/psutil/_pswindows.py new file mode 100644 index 0000000000000000000000000000000000000000..3ee1eafce0bb08f86df6d08d30a66bb2058c6f7c --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/psutil/_pswindows.py @@ -0,0 +1,1096 @@ +# Copyright (c) 2009, Giampaolo Rodola'. All rights reserved. +# Use of this source code is governed by a BSD-style license that can be +# found in the LICENSE file. + +"""Windows platform implementation.""" + +import contextlib +import enum +import functools +import os +import signal +import sys +import threading +import time + +from . import _common +from . import _ntuples as ntp +from ._common import ENCODING +from ._common import AccessDenied +from ._common import NoSuchProcess +from ._common import TimeoutExpired +from ._common import conn_tmap +from ._common import conn_to_ntuple +from ._common import debug +from ._common import isfile_strict +from ._common import memoize +from ._common import memoize_when_activated +from ._common import parse_environ_block +from ._common import usage_percent +from ._psutil_windows import ABOVE_NORMAL_PRIORITY_CLASS +from ._psutil_windows import BELOW_NORMAL_PRIORITY_CLASS +from ._psutil_windows import HIGH_PRIORITY_CLASS +from ._psutil_windows import IDLE_PRIORITY_CLASS +from ._psutil_windows import NORMAL_PRIORITY_CLASS +from ._psutil_windows import REALTIME_PRIORITY_CLASS + +try: + from . import _psutil_windows as cext +except ImportError as err: + if ( + str(err).lower().startswith("dll load failed") + and sys.getwindowsversion()[0] < 6 + ): + # We may get here if: + # 1) we are on an old Windows version + # 2) psutil was installed via pip + wheel + # See: https://github.com/giampaolo/psutil/issues/811 + msg = "this Windows version is too old (< Windows Vista); " + msg += "psutil 3.4.2 is the latest version which supports Windows " + msg += "2000, XP and 2003 server" + raise RuntimeError(msg) from err + else: + raise + + +# process priority constants, import from __init__.py: +# http://msdn.microsoft.com/en-us/library/ms686219(v=vs.85).aspx +# fmt: off +__extra__all__ = [ + "win_service_iter", "win_service_get", + # Process priority + "ABOVE_NORMAL_PRIORITY_CLASS", "BELOW_NORMAL_PRIORITY_CLASS", + "HIGH_PRIORITY_CLASS", "IDLE_PRIORITY_CLASS", "NORMAL_PRIORITY_CLASS", + "REALTIME_PRIORITY_CLASS", + # IO priority + "IOPRIO_VERYLOW", "IOPRIO_LOW", "IOPRIO_NORMAL", "IOPRIO_HIGH", + # others + "CONN_DELETE_TCB", "AF_LINK", +] +# fmt: on + + +# ===================================================================== +# --- globals +# ===================================================================== + +CONN_DELETE_TCB = "DELETE_TCB" +ERROR_PARTIAL_COPY = 299 +PYPY = '__pypy__' in sys.builtin_module_names + +AddressFamily = enum.IntEnum('AddressFamily', {'AF_LINK': -1}) +AF_LINK = AddressFamily.AF_LINK + +TCP_STATUSES = { + cext.MIB_TCP_STATE_ESTAB: _common.CONN_ESTABLISHED, + cext.MIB_TCP_STATE_SYN_SENT: _common.CONN_SYN_SENT, + cext.MIB_TCP_STATE_SYN_RCVD: _common.CONN_SYN_RECV, + cext.MIB_TCP_STATE_FIN_WAIT1: _common.CONN_FIN_WAIT1, + cext.MIB_TCP_STATE_FIN_WAIT2: _common.CONN_FIN_WAIT2, + cext.MIB_TCP_STATE_TIME_WAIT: _common.CONN_TIME_WAIT, + cext.MIB_TCP_STATE_CLOSED: _common.CONN_CLOSE, + cext.MIB_TCP_STATE_CLOSE_WAIT: _common.CONN_CLOSE_WAIT, + cext.MIB_TCP_STATE_LAST_ACK: _common.CONN_LAST_ACK, + cext.MIB_TCP_STATE_LISTEN: _common.CONN_LISTEN, + cext.MIB_TCP_STATE_CLOSING: _common.CONN_CLOSING, + cext.MIB_TCP_STATE_DELETE_TCB: CONN_DELETE_TCB, + cext.PSUTIL_CONN_NONE: _common.CONN_NONE, +} + + +class Priority(enum.IntEnum): + ABOVE_NORMAL_PRIORITY_CLASS = ABOVE_NORMAL_PRIORITY_CLASS + BELOW_NORMAL_PRIORITY_CLASS = BELOW_NORMAL_PRIORITY_CLASS + HIGH_PRIORITY_CLASS = HIGH_PRIORITY_CLASS + IDLE_PRIORITY_CLASS = IDLE_PRIORITY_CLASS + NORMAL_PRIORITY_CLASS = NORMAL_PRIORITY_CLASS + REALTIME_PRIORITY_CLASS = REALTIME_PRIORITY_CLASS + + +globals().update(Priority.__members__) + + +class IOPriority(enum.IntEnum): + IOPRIO_VERYLOW = 0 + IOPRIO_LOW = 1 + IOPRIO_NORMAL = 2 + IOPRIO_HIGH = 3 + + +globals().update(IOPriority.__members__) + +pinfo_map = dict( + num_handles=0, + ctx_switches=1, + user_time=2, + kernel_time=3, + create_time=4, + num_threads=5, + io_rcount=6, + io_wcount=7, + io_rbytes=8, + io_wbytes=9, + io_count_others=10, + io_bytes_others=11, + num_page_faults=12, + peak_wset=13, + wset=14, + peak_paged_pool=15, + paged_pool=16, + peak_non_paged_pool=17, + non_paged_pool=18, + pagefile=19, + peak_pagefile=20, + mem_private=21, +) + + +# ===================================================================== +# --- utils +# ===================================================================== + + +@functools.lru_cache(maxsize=512) +def convert_dos_path(s): + r"""Convert paths using native DOS format like: + "\Device\HarddiskVolume1\Windows\systemew\file.txt" or + "\??\C:\Windows\systemew\file.txt" + into: + "C:\Windows\systemew\file.txt". + """ + if s.startswith('\\\\'): + return s + rawdrive = '\\'.join(s.split('\\')[:3]) + if rawdrive in {"\\??\\UNC", "\\Device\\Mup"}: + rawdrive = '\\'.join(s.split('\\')[:5]) + driveletter = '\\\\' + '\\'.join(s.split('\\')[3:5]) + elif rawdrive.startswith('\\??\\'): + driveletter = s.split('\\')[2] + else: + driveletter = cext.QueryDosDevice(rawdrive) + remainder = s[len(rawdrive) :] + return os.path.join(driveletter, remainder) + + +@memoize +def getpagesize(): + return cext.getpagesize() + + +# ===================================================================== +# --- memory +# ===================================================================== + + +def virtual_memory(): + """System virtual memory as a namedtuple.""" + mem = cext.virtual_mem() + totphys, availphys, _totsys, _availsys = mem + total = totphys + avail = availphys + free = availphys + used = total - avail + percent = usage_percent((total - avail), total, round_=1) + return ntp.svmem(total, avail, percent, used, free) + + +def swap_memory(): + """Swap system memory as a (total, used, free, sin, sout) tuple.""" + mem = cext.virtual_mem() + + total_phys = mem[0] + total_system = mem[2] + + # system memory (commit total/limit) is the sum of physical and swap + # thus physical memory values need to be subtracted to get swap values + total = total_system - total_phys + # commit total is incremented immediately (decrementing free_system) + # while the corresponding free physical value is not decremented until + # pages are accessed, so we can't use free system memory for swap. + # instead, we calculate page file usage based on performance counter + if total > 0: + percentswap = cext.swap_percent() + used = int(0.01 * percentswap * total) + else: + percentswap = 0.0 + used = 0 + + free = total - used + percent = round(percentswap, 1) + return ntp.sswap(total, used, free, percent, 0, 0) + + +# malloc / heap functions +heap_info = cext.heap_info +heap_trim = cext.heap_trim + + +# ===================================================================== +# --- disk +# ===================================================================== + + +disk_io_counters = cext.disk_io_counters + + +def disk_usage(path): + """Return disk usage associated with path.""" + if isinstance(path, bytes): + # XXX: do we want to use "strict"? Probably yes, in order + # to fail immediately. After all we are accepting input here... + path = path.decode(ENCODING, errors="strict") + total, used, free = cext.disk_usage(path) + percent = usage_percent(used, total, round_=1) + return ntp.sdiskusage(total, used, free, percent) + + +def disk_partitions(all): + """Return disk partitions.""" + rawlist = cext.disk_partitions(all) + return [ntp.sdiskpart(*x) for x in rawlist] + + +# ===================================================================== +# --- CPU +# ===================================================================== + + +def cpu_times(): + """Return system CPU times as a named tuple.""" + user, system, idle = cext.cpu_times() + # Internally, GetSystemTimes() is used, and it doesn't return + # interrupt and dpc times. cext.per_cpu_times() does, so we + # rely on it to get those only. + percpu_summed = ntp.scputimes( + *[sum(n) for n in zip(*cext.per_cpu_times())] + ) + return ntp.scputimes( + user, system, idle, percpu_summed.interrupt, percpu_summed.dpc + ) + + +def per_cpu_times(): + """Return system per-CPU times as a list of named tuples.""" + ret = [] + for user, system, idle, interrupt, dpc in cext.per_cpu_times(): + item = ntp.scputimes(user, system, idle, interrupt, dpc) + ret.append(item) + return ret + + +def cpu_count_logical(): + """Return the number of logical CPUs in the system.""" + return cext.cpu_count_logical() + + +def cpu_count_cores(): + """Return the number of CPU cores in the system.""" + return cext.cpu_count_cores() + + +def cpu_stats(): + """Return CPU statistics.""" + ctx_switches, interrupts, _dpcs, syscalls = cext.cpu_stats() + soft_interrupts = 0 + return ntp.scpustats(ctx_switches, interrupts, soft_interrupts, syscalls) + + +def cpu_freq(): + """Return CPU frequency. + On Windows per-cpu frequency is not supported. + """ + curr, max_ = cext.cpu_freq() + min_ = 0.0 + return [ntp.scpufreq(float(curr), min_, float(max_))] + + +_loadavg_initialized = False +_lock = threading.Lock() + + +def _getloadavg_impl(): + # Drop to 2 decimal points which is what Linux does + raw_loads = cext.getloadavg() + return tuple(round(load, 2) for load in raw_loads) + + +def getloadavg(): + """Return the number of processes in the system run queue averaged + over the last 1, 5, and 15 minutes respectively as a tuple. + """ + global _loadavg_initialized + + if _loadavg_initialized: + return _getloadavg_impl() + + with _lock: + if not _loadavg_initialized: + cext.init_loadavg_counter() + _loadavg_initialized = True + + return _getloadavg_impl() + + +# ===================================================================== +# --- network +# ===================================================================== + + +def net_connections(kind, _pid=-1): + """Return socket connections. If pid == -1 return system-wide + connections (as opposed to connections opened by one process only). + """ + families, types = conn_tmap[kind] + rawlist = cext.net_connections(_pid, families, types) + ret = set() + for item in rawlist: + fd, fam, type, laddr, raddr, status, pid = item + nt = conn_to_ntuple( + fd, + fam, + type, + laddr, + raddr, + status, + TCP_STATUSES, + pid=pid if _pid == -1 else None, + ) + ret.add(nt) + return list(ret) + + +def net_if_stats(): + """Get NIC stats (isup, duplex, speed, mtu).""" + ret = {} + rawdict = cext.net_if_stats() + for name, items in rawdict.items(): + isup, duplex, speed, mtu = items + if hasattr(_common, 'NicDuplex'): + duplex = _common.NicDuplex(duplex) + ret[name] = ntp.snicstats(isup, duplex, speed, mtu, '') + return ret + + +def net_io_counters(): + """Return network I/O statistics for every network interface + installed on the system as a dict of raw tuples. + """ + return cext.net_io_counters() + + +def net_if_addrs(): + """Return the addresses associated to each NIC.""" + return cext.net_if_addrs() + + +# ===================================================================== +# --- sensors +# ===================================================================== + + +def sensors_battery(): + """Return battery information.""" + # For constants meaning see: + # https://msdn.microsoft.com/en-us/library/windows/desktop/ + # aa373232(v=vs.85).aspx + acline_status, flags, percent, secsleft = cext.sensors_battery() + power_plugged = acline_status == 1 + no_battery = bool(flags & 128) + charging = bool(flags & 8) + + if no_battery: + return None + if power_plugged or charging: + secsleft = _common.POWER_TIME_UNLIMITED + elif secsleft == -1: + secsleft = _common.POWER_TIME_UNKNOWN + + return ntp.sbattery(percent, secsleft, power_plugged) + + +# ===================================================================== +# --- other system functions +# ===================================================================== + + +_last_btime = 0 + + +def boot_time(): + """The system boot time expressed in seconds since the epoch. This + also includes the time spent during hybernate / suspend. + """ + # This dirty hack is to adjust the precision of the returned + # value which may have a 1 second fluctuation, see: + # https://github.com/giampaolo/psutil/issues/1007 + global _last_btime + ret = time.time() - cext.uptime() + if abs(ret - _last_btime) <= 1: + return _last_btime + else: + _last_btime = ret + return ret + + +def users(): + """Return currently connected users as a list of namedtuples.""" + retlist = [] + rawlist = cext.users() + for item in rawlist: + user, hostname, tstamp = item + nt = ntp.suser(user, None, hostname, tstamp, None) + retlist.append(nt) + return retlist + + +# ===================================================================== +# --- Windows services +# ===================================================================== + + +def win_service_iter(): + """Yields a list of WindowsService instances.""" + for name, display_name in cext.winservice_enumerate(): + yield WindowsService(name, display_name) + + +def win_service_get(name): + """Open a Windows service and return it as a WindowsService instance.""" + service = WindowsService(name, None) + service._display_name = service._query_config()['display_name'] + return service + + +class WindowsService: # noqa: PLW1641 + """Represents an installed Windows service.""" + + def __init__(self, name, display_name): + self._name = name + self._display_name = display_name + + def __str__(self): + details = f"(name={self._name!r}, display_name={self._display_name!r})" + return f"{self.__class__.__name__}{details}" + + def __repr__(self): + return f"<{self.__str__()} at {id(self)}>" + + def __eq__(self, other): + # Test for equality with another WindosService object based + # on name. + if not isinstance(other, WindowsService): + return NotImplemented + return self._name == other._name + + def __ne__(self, other): + return not self == other + + def _query_config(self): + with self._wrap_exceptions(): + display_name, binpath, username, start_type = ( + cext.winservice_query_config(self._name) + ) + # XXX - update _self.display_name? + return dict( + display_name=display_name, + binpath=binpath, + username=username, + start_type=start_type, + ) + + def _query_status(self): + with self._wrap_exceptions(): + status, pid = cext.winservice_query_status(self._name) + if pid == 0: + pid = None + return dict(status=status, pid=pid) + + @contextlib.contextmanager + def _wrap_exceptions(self): + """Ctx manager which translates bare OSError and WindowsError + exceptions into NoSuchProcess and AccessDenied. + """ + try: + yield + except OSError as err: + name = self._name + if is_permission_err(err): + msg = ( + f"service {name!r} is not querable (not enough privileges)" + ) + raise AccessDenied(pid=None, name=name, msg=msg) from err + elif err.winerror in { + cext.ERROR_INVALID_NAME, + cext.ERROR_SERVICE_DOES_NOT_EXIST, + }: + msg = f"service {name!r} does not exist" + raise NoSuchProcess(pid=None, name=name, msg=msg) from err + else: + raise + + # config query + + def name(self): + """The service name. This string is how a service is referenced + and can be passed to win_service_get() to get a new + WindowsService instance. + """ + return self._name + + def display_name(self): + """The service display name. The value is cached when this class + is instantiated. + """ + return self._display_name + + def binpath(self): + """The fully qualified path to the service binary/exe file as + a string, including command line arguments. + """ + return self._query_config()['binpath'] + + def username(self): + """The name of the user that owns this service.""" + return self._query_config()['username'] + + def start_type(self): + """A string which can either be "automatic", "manual" or + "disabled". + """ + return self._query_config()['start_type'] + + # status query + + def pid(self): + """The process PID, if any, else None. This can be passed + to Process class to control the service's process. + """ + return self._query_status()['pid'] + + def status(self): + """Service status as a string.""" + return self._query_status()['status'] + + def description(self): + """Service long description.""" + return cext.winservice_query_descr(self.name()) + + # utils + + def as_dict(self): + """Utility method retrieving all the information above as a + dictionary. + """ + d = self._query_config() + d.update(self._query_status()) + d['name'] = self.name() + d['display_name'] = self.display_name() + d['description'] = self.description() + return d + + # actions + # XXX: the necessary C bindings for start() and stop() are + # implemented but for now I prefer not to expose them. + # I may change my mind in the future. Reasons: + # - they require Administrator privileges + # - can't implement a timeout for stop() (unless by using a thread, + # which sucks) + # - would require adding ServiceAlreadyStarted and + # ServiceAlreadyStopped exceptions, adding two new APIs. + # - we might also want to have modify(), which would basically mean + # rewriting win32serviceutil.ChangeServiceConfig, which involves a + # lot of stuff (and API constants which would pollute the API), see: + # http://pyxr.sourceforge.net/PyXR/c/python24/lib/site-packages/ + # win32/lib/win32serviceutil.py.html#0175 + # - psutil is typically about "read only" monitoring stuff; + # win_service_* APIs should only be used to retrieve a service and + # check whether it's running + + # def start(self, timeout=None): + # with self._wrap_exceptions(): + # cext.winservice_start(self.name()) + # if timeout: + # giveup_at = time.time() + timeout + # while True: + # if self.status() == "running": + # return + # else: + # if time.time() > giveup_at: + # raise TimeoutExpired(timeout) + # else: + # time.sleep(.1) + + # def stop(self): + # # Note: timeout is not implemented because it's just not + # # possible, see: + # # http://stackoverflow.com/questions/11973228/ + # with self._wrap_exceptions(): + # return cext.winservice_stop(self.name()) + + +# ===================================================================== +# --- processes +# ===================================================================== + + +pids = cext.pids +pid_exists = cext.pid_exists +ppid_map = cext.ppid_map # used internally by Process.children() + + +def is_permission_err(exc): + """Return True if this is a permission error.""" + assert isinstance(exc, OSError), exc + return isinstance(exc, PermissionError) or exc.winerror in { + cext.ERROR_ACCESS_DENIED, + cext.ERROR_PRIVILEGE_NOT_HELD, + } + + +def convert_oserror(exc, pid=None, name=None): + """Convert OSError into NoSuchProcess or AccessDenied.""" + assert isinstance(exc, OSError), exc + if is_permission_err(exc): + return AccessDenied(pid=pid, name=name) + if isinstance(exc, ProcessLookupError): + return NoSuchProcess(pid=pid, name=name) + raise exc + + +def wrap_exceptions(fun): + """Decorator which converts OSError into NoSuchProcess or AccessDenied.""" + + @functools.wraps(fun) + def wrapper(self, *args, **kwargs): + try: + return fun(self, *args, **kwargs) + except OSError as err: + raise convert_oserror(err, pid=self.pid, name=self._name) from err + + return wrapper + + +def retry_error_partial_copy(fun): + """Workaround for https://github.com/giampaolo/psutil/issues/875. + See: https://stackoverflow.com/questions/4457745#4457745. + """ + + @functools.wraps(fun) + def wrapper(self, *args, **kwargs): + delay = 0.0001 + times = 33 + for _ in range(times): # retries for roughly 1 second + try: + return fun(self, *args, **kwargs) + except OSError as _: + err = _ + if err.winerror == ERROR_PARTIAL_COPY: + time.sleep(delay) + delay = min(delay * 2, 0.04) + continue + raise + msg = ( + f"{fun} retried {times} times, converted to AccessDenied as it's " + f"still returning {err}" + ) + raise AccessDenied(pid=self.pid, name=self._name, msg=msg) + + return wrapper + + +class Process: + """Wrapper class around underlying C implementation.""" + + __slots__ = ["_cache", "_name", "_ppid", "pid"] + + def __init__(self, pid): + self.pid = pid + self._name = None + self._ppid = None + + # --- oneshot() stuff + + def oneshot_enter(self): + self._proc_info.cache_activate(self) + self.exe.cache_activate(self) + + def oneshot_exit(self): + self._proc_info.cache_deactivate(self) + self.exe.cache_deactivate(self) + + @memoize_when_activated + def _proc_info(self): + """Return multiple information about this process as a + raw tuple. + """ + ret = cext.proc_info(self.pid) + assert len(ret) == len(pinfo_map) + return ret + + def name(self): + """Return process name, which on Windows is always the final + part of the executable. + """ + # This is how PIDs 0 and 4 are always represented in taskmgr + # and process-hacker. + if self.pid == 0: + return "System Idle Process" + if self.pid == 4: + return "System" + return os.path.basename(self.exe()) + + @wrap_exceptions + @memoize_when_activated + def exe(self): + if PYPY: + try: + exe = cext.proc_exe(self.pid) + except OSError as err: + # 24 = ERROR_TOO_MANY_OPEN_FILES. Not sure why this happens + # (perhaps PyPy's JIT delaying garbage collection of files?). + if err.errno == 24: + debug(f"{err!r} translated into AccessDenied") + raise AccessDenied(self.pid, self._name) from err + raise + else: + exe = cext.proc_exe(self.pid) + if exe.startswith('\\'): + return convert_dos_path(exe) + return exe # May be "Registry", "MemCompression", ... + + @wrap_exceptions + @retry_error_partial_copy + def cmdline(self): + if cext.WINVER >= cext.WINDOWS_8_1: + # PEB method detects cmdline changes but requires more + # privileges: https://github.com/giampaolo/psutil/pull/1398 + try: + return cext.proc_cmdline(self.pid, use_peb=True) + except OSError as err: + if is_permission_err(err): + return cext.proc_cmdline(self.pid, use_peb=False) + else: + raise + else: + return cext.proc_cmdline(self.pid, use_peb=True) + + @wrap_exceptions + @retry_error_partial_copy + def environ(self): + s = cext.proc_environ(self.pid) + return parse_environ_block(s) + + def ppid(self): + try: + return ppid_map()[self.pid] + except KeyError: + raise NoSuchProcess(self.pid, self._name) from None + + def _get_raw_meminfo(self): + try: + return cext.proc_memory_info(self.pid) + except OSError as err: + if is_permission_err(err): + # TODO: the C ext can probably be refactored in order + # to get this from cext.proc_info() + debug("attempting memory_info() fallback (slower)") + info = self._proc_info() + return ( + info[pinfo_map['num_page_faults']], + info[pinfo_map['peak_wset']], + info[pinfo_map['wset']], + info[pinfo_map['peak_paged_pool']], + info[pinfo_map['paged_pool']], + info[pinfo_map['peak_non_paged_pool']], + info[pinfo_map['non_paged_pool']], + info[pinfo_map['pagefile']], + info[pinfo_map['peak_pagefile']], + info[pinfo_map['mem_private']], + ) + raise + + @wrap_exceptions + def memory_info(self): + # on Windows RSS == WorkingSetSize and VSM == PagefileUsage. + # Underlying C function returns fields of PROCESS_MEMORY_COUNTERS + # struct. + t = self._get_raw_meminfo() + rss = t[2] # wset + vms = t[7] # pagefile + return ntp.pmem(*(rss, vms) + t) + + @wrap_exceptions + def memory_full_info(self): + basic_mem = self.memory_info() + uss = cext.proc_memory_uss(self.pid) + uss *= getpagesize() + return ntp.pfullmem(*basic_mem + (uss,)) + + def memory_maps(self): + try: + raw = cext.proc_memory_maps(self.pid) + except OSError as err: + # XXX - can't use wrap_exceptions decorator as we're + # returning a generator; probably needs refactoring. + raise convert_oserror(err, self.pid, self._name) from err + else: + for addr, perm, path, rss in raw: + path = convert_dos_path(path) + addr = hex(addr) + yield (addr, perm, path, rss) + + @wrap_exceptions + def kill(self): + return cext.proc_kill(self.pid) + + @wrap_exceptions + def send_signal(self, sig): + if sig == signal.SIGTERM: + cext.proc_kill(self.pid) + elif sig in {signal.CTRL_C_EVENT, signal.CTRL_BREAK_EVENT}: + os.kill(self.pid, sig) + else: + msg = ( + "only SIGTERM, CTRL_C_EVENT and CTRL_BREAK_EVENT signals " + "are supported on Windows" + ) + raise ValueError(msg) + + @wrap_exceptions + def wait(self, timeout=None): + if timeout is None: + cext_timeout = cext.INFINITE + else: + # WaitForSingleObject() expects time in milliseconds. + cext_timeout = int(timeout * 1000) + + timer = getattr(time, 'monotonic', time.time) + stop_at = timer() + timeout if timeout is not None else None + + try: + # Exit code is supposed to come from GetExitCodeProcess(). + # May also be None if OpenProcess() failed with + # ERROR_INVALID_PARAMETER, meaning PID is already gone. + exit_code = cext.proc_wait(self.pid, cext_timeout) + except cext.TimeoutExpired as err: + # WaitForSingleObject() returned WAIT_TIMEOUT. Just raise. + raise TimeoutExpired(timeout, self.pid, self._name) from err + except cext.TimeoutAbandoned: + # WaitForSingleObject() returned WAIT_ABANDONED, see: + # https://github.com/giampaolo/psutil/issues/1224 + # We'll just rely on the internal polling and return None + # when the PID disappears. Subprocess module does the same + # (return None): + # https://github.com/python/cpython/blob/ + # be50a7b627d0aa37e08fa8e2d5568891f19903ce/ + # Lib/subprocess.py#L1193-L1194 + exit_code = None + + # At this point WaitForSingleObject() returned WAIT_OBJECT_0, + # meaning the process is gone. Stupidly there are cases where + # its PID may still stick around so we do a further internal + # polling. + delay = 0.0001 + while True: + if not pid_exists(self.pid): + return exit_code + if stop_at and timer() >= stop_at: + raise TimeoutExpired(timeout, pid=self.pid, name=self._name) + time.sleep(delay) + delay = min(delay * 2, 0.04) # incremental delay + + @wrap_exceptions + def username(self): + if self.pid in {0, 4}: + return 'NT AUTHORITY\\SYSTEM' + domain, user = cext.proc_username(self.pid) + return f"{domain}\\{user}" + + @wrap_exceptions + def create_time(self, fast_only=False): + # Note: proc_times() not put under oneshot() 'cause create_time() + # is already cached by the main Process class. + try: + _user, _system, created = cext.proc_times(self.pid) + return created + except OSError as err: + if is_permission_err(err): + if fast_only: + raise + debug("attempting create_time() fallback (slower)") + return self._proc_info()[pinfo_map['create_time']] + raise + + @wrap_exceptions + def num_threads(self): + return self._proc_info()[pinfo_map['num_threads']] + + @wrap_exceptions + def threads(self): + rawlist = cext.proc_threads(self.pid) + retlist = [] + for thread_id, utime, stime in rawlist: + ntuple = ntp.pthread(thread_id, utime, stime) + retlist.append(ntuple) + return retlist + + @wrap_exceptions + def cpu_times(self): + try: + user, system, _created = cext.proc_times(self.pid) + except OSError as err: + if not is_permission_err(err): + raise + debug("attempting cpu_times() fallback (slower)") + info = self._proc_info() + user = info[pinfo_map['user_time']] + system = info[pinfo_map['kernel_time']] + # Children user/system times are not retrievable (set to 0). + return ntp.pcputimes(user, system, 0.0, 0.0) + + @wrap_exceptions + def suspend(self): + cext.proc_suspend_or_resume(self.pid, True) + + @wrap_exceptions + def resume(self): + cext.proc_suspend_or_resume(self.pid, False) + + @wrap_exceptions + @retry_error_partial_copy + def cwd(self): + if self.pid in {0, 4}: + raise AccessDenied(self.pid, self._name) + # return a normalized pathname since the native C function appends + # "\\" at the and of the path + path = cext.proc_cwd(self.pid) + return os.path.normpath(path) + + @wrap_exceptions + def open_files(self): + if self.pid in {0, 4}: + return [] + ret = set() + # Filenames come in in native format like: + # "\Device\HarddiskVolume1\Windows\systemew\file.txt" + # Convert the first part in the corresponding drive letter + # (e.g. "C:\") by using Windows's QueryDosDevice() + raw_file_names = cext.proc_open_files(self.pid) + for file in raw_file_names: + file = convert_dos_path(file) + if isfile_strict(file): + ntuple = ntp.popenfile(file, -1) + ret.add(ntuple) + return list(ret) + + @wrap_exceptions + def net_connections(self, kind='inet'): + return net_connections(kind, _pid=self.pid) + + @wrap_exceptions + def nice_get(self): + value = cext.proc_priority_get(self.pid) + value = Priority(value) + return value + + @wrap_exceptions + def nice_set(self, value): + return cext.proc_priority_set(self.pid, value) + + @wrap_exceptions + def ionice_get(self): + ret = cext.proc_io_priority_get(self.pid) + ret = IOPriority(ret) + return ret + + @wrap_exceptions + def ionice_set(self, ioclass, value): + if value: + msg = "value argument not accepted on Windows" + raise TypeError(msg) + if ioclass not in { + IOPriority.IOPRIO_VERYLOW, + IOPriority.IOPRIO_LOW, + IOPriority.IOPRIO_NORMAL, + IOPriority.IOPRIO_HIGH, + }: + msg = f"{ioclass} is not a valid priority" + raise ValueError(msg) + cext.proc_io_priority_set(self.pid, ioclass) + + @wrap_exceptions + def io_counters(self): + try: + ret = cext.proc_io_counters(self.pid) + except OSError as err: + if not is_permission_err(err): + raise + debug("attempting io_counters() fallback (slower)") + info = self._proc_info() + ret = ( + info[pinfo_map['io_rcount']], + info[pinfo_map['io_wcount']], + info[pinfo_map['io_rbytes']], + info[pinfo_map['io_wbytes']], + info[pinfo_map['io_count_others']], + info[pinfo_map['io_bytes_others']], + ) + return ntp.pio(*ret) + + @wrap_exceptions + def status(self): + suspended = cext.proc_is_suspended(self.pid) + if suspended: + return _common.STATUS_STOPPED + else: + return _common.STATUS_RUNNING + + @wrap_exceptions + def cpu_affinity_get(self): + def from_bitmask(x): + return [i for i in range(64) if (1 << i) & x] + + bitmask = cext.proc_cpu_affinity_get(self.pid) + return from_bitmask(bitmask) + + @wrap_exceptions + def cpu_affinity_set(self, value): + def to_bitmask(ls): + if not ls: + msg = f"invalid argument {ls!r}" + raise ValueError(msg) + out = 0 + for b in ls: + out |= 2**b + return out + + # SetProcessAffinityMask() states that ERROR_INVALID_PARAMETER + # is returned for an invalid CPU but this seems not to be true, + # therefore we check CPUs validy beforehand. + allcpus = list(range(len(per_cpu_times()))) + for cpu in value: + if cpu not in allcpus: + if not isinstance(cpu, int): + msg = f"invalid CPU {cpu!r}; an integer is required" + raise TypeError(msg) + msg = f"invalid CPU {cpu!r}" + raise ValueError(msg) + + bitmask = to_bitmask(value) + cext.proc_cpu_affinity_set(self.pid, bitmask) + + @wrap_exceptions + def num_handles(self): + try: + return cext.proc_num_handles(self.pid) + except OSError as err: + if is_permission_err(err): + debug("attempting num_handles() fallback (slower)") + return self._proc_info()[pinfo_map['num_handles']] + raise + + @wrap_exceptions + def num_ctx_switches(self): + ctx_switches = self._proc_info()[pinfo_map['ctx_switches']] + # only voluntary ctx switches are supported + return ntp.pctxsw(ctx_switches, 0) diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/pytablewriter-1.2.1.dist-info/INSTALLER b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/pytablewriter-1.2.1.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/pytablewriter-1.2.1.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/pytablewriter-1.2.1.dist-info/LICENSE b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/pytablewriter-1.2.1.dist-info/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..3e3c2f9d5e2bc9c160cda612c58f5fb61df672d0 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/pytablewriter-1.2.1.dist-info/LICENSE @@ -0,0 +1,21 @@ +The MIT License (MIT) + +Copyright (c) 2016-2025 Tsuyoshi Hombashi + +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. diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/pytablewriter-1.2.1.dist-info/METADATA b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/pytablewriter-1.2.1.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..013bef3e8bd46cf7b5c6c3b40ec30a4553b338ce --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/pytablewriter-1.2.1.dist-info/METADATA @@ -0,0 +1,932 @@ +Metadata-Version: 2.1 +Name: pytablewriter +Version: 1.2.1 +Summary: pytablewriter is a Python library to write a table in various formats: AsciiDoc / CSV / Elasticsearch / HTML / JavaScript / JSON / LaTeX / LDJSON / LTSV / Markdown / MediaWiki / NumPy / Excel / Pandas / Python / reStructuredText / SQLite / TOML / TSV / YAML. +Home-page: https://github.com/thombashi/pytablewriter +Author: Tsuyoshi Hombashi +Author-email: tsuyoshi.hombashi@gmail.com +License: MIT License +Project-URL: Changelog, https://github.com/thombashi/pytablewriter/blob/master/CHANGELOG.md +Project-URL: Documentation, https://pytablewriter.rtfd.io/ +Project-URL: Funding, https://github.com/sponsors/thombashi +Project-URL: Source, https://github.com/thombashi/pytablewriter +Project-URL: Tracker, https://github.com/thombashi/pytablewriter/issues +Keywords: AsciiDoc,table,CSV,Excel,JavaScript,JSON,LaTeX,LTSV,Markdown,MediaWiki,HTML,pandas,reStructuredText,SQLite,TSV,TOML +Classifier: Development Status :: 5 - Production/Stable +Classifier: Intended Audience :: Developers +Classifier: Intended Audience :: Information Technology +Classifier: License :: OSI Approved :: MIT License +Classifier: Operating System :: OS Independent +Classifier: Programming Language :: Python :: 3 +Classifier: Programming Language :: Python :: 3.9 +Classifier: Programming Language :: Python :: 3.10 +Classifier: Programming Language :: Python :: 3.11 +Classifier: Programming Language :: Python :: 3.12 +Classifier: Programming Language :: Python :: 3.13 +Classifier: Programming Language :: Python :: 3 :: Only +Classifier: Programming Language :: Python :: Implementation :: CPython +Classifier: Programming Language :: Python :: Implementation :: PyPy +Classifier: Topic :: Software Development :: Code Generators +Classifier: Topic :: Software Development :: Libraries +Classifier: Topic :: Software Development :: Libraries :: Python Modules +Classifier: Topic :: Text Processing +Classifier: Topic :: Text Processing :: Markup :: HTML +Classifier: Topic :: Text Processing :: Markup :: LaTeX +Classifier: Topic :: Text Processing :: Markup :: Markdown +Classifier: Topic :: Text Processing :: Markup :: reStructuredText +Classifier: Typing :: Typed +Requires-Python: >=3.9 +Description-Content-Type: text/x-rst +License-File: LICENSE +Requires-Dist: setuptools>=38.3.0 +Requires-Dist: DataProperty<2,>=1.1.0 +Requires-Dist: mbstrdecoder<2,>=1.0.0 +Requires-Dist: pathvalidate<4,>=2.3.0 +Requires-Dist: tabledata<2,>=1.3.1 +Requires-Dist: tcolorpy<1,>=0.0.5 +Requires-Dist: typepy[datetime]<2,>=1.3.2 +Provides-Extra: all +Requires-Dist: xlwt; extra == "all" +Requires-Dist: XlsxWriter<4,>=0.9.6; extra == "all" +Requires-Dist: elasticsearch<9,>=8.0.1; extra == "all" +Requires-Dist: pytablereader<2,>=0.31.3; extra == "all" +Requires-Dist: dominate<3,>=2.1.5; extra == "all" +Requires-Dist: loguru<1,>=0.4.1; extra == "all" +Requires-Dist: SimpleSQLite<2,>=1.3.2; extra == "all" +Requires-Dist: pytablewriter-altrow-theme<1,>=0.2.0; extra == "all" +Requires-Dist: pytablewriter-altcol-theme<1,>=0.1.0; extra == "all" +Requires-Dist: toml<1,>=0.9.3; extra == "all" +Requires-Dist: PyYAML<7,>=3.11; extra == "all" +Requires-Dist: simplejson<4,>=3.8.1; extra == "all" +Requires-Dist: pandas<3,>=0.25.3; extra == "all" +Provides-Extra: docs +Requires-Dist: sphinx_rtd_theme>=1.2.2; extra == "docs" +Requires-Dist: Sphinx>=2.4; extra == "docs" +Requires-Dist: xlwt; extra == "docs" +Requires-Dist: XlsxWriter<4,>=0.9.6; extra == "docs" +Requires-Dist: elasticsearch<9,>=8.0.1; extra == "docs" +Requires-Dist: pytablereader<2,>=0.31.3; extra == "docs" +Requires-Dist: dominate<3,>=2.1.5; extra == "docs" +Requires-Dist: loguru<1,>=0.4.1; extra == "docs" +Requires-Dist: SimpleSQLite<2,>=1.3.2; extra == "docs" +Requires-Dist: pytablewriter-altrow-theme<1,>=0.2.0; extra == "docs" +Requires-Dist: pytablewriter-altcol-theme<1,>=0.1.0; extra == "docs" +Requires-Dist: toml<1,>=0.9.3; extra == "docs" +Requires-Dist: PyYAML<7,>=3.11; extra == "docs" +Requires-Dist: simplejson<4,>=3.8.1; extra == "docs" +Requires-Dist: pandas<3,>=0.25.3; extra == "docs" +Provides-Extra: es +Requires-Dist: elasticsearch<9,>=8.0.1; extra == "es" +Provides-Extra: es8 +Requires-Dist: elasticsearch<9,>=8.0.1; extra == "es8" +Provides-Extra: excel +Requires-Dist: xlwt; extra == "excel" +Requires-Dist: XlsxWriter<4,>=0.9.6; extra == "excel" +Provides-Extra: from +Requires-Dist: pytablereader<2,>=0.31.3; extra == "from" +Provides-Extra: html +Requires-Dist: dominate<3,>=2.1.5; extra == "html" +Provides-Extra: logging +Requires-Dist: loguru<1,>=0.4.1; extra == "logging" +Provides-Extra: pandas +Requires-Dist: pandas<3,>=0.25.3; extra == "pandas" +Provides-Extra: sqlite +Requires-Dist: SimpleSQLite<2,>=1.3.2; extra == "sqlite" +Provides-Extra: test +Requires-Dist: pytest-md-report>=0.6.2; extra == "test" +Requires-Dist: pytablewriter-altcol-theme<1,>=0.1.0; extra == "test" +Requires-Dist: PyYAML<7,>=3.11; extra == "test" +Requires-Dist: pytablereader[excel,sqlite]>=0.31.3; extra == "test" +Requires-Dist: dominate<3,>=2.1.5; extra == "test" +Requires-Dist: pytablewriter-altrow-theme<1,>=0.2.0; extra == "test" +Requires-Dist: toml<1,>=0.9.3; extra == "test" +Requires-Dist: simplejson<4,>=3.8.1; extra == "test" +Requires-Dist: pytest>=6.0.1; extra == "test" +Requires-Dist: beautifulsoup4>=4.10; extra == "test" +Requires-Dist: SimpleSQLite<2,>=1.3.2; extra == "test" +Requires-Dist: pytablereader<2,>=0.31.3; extra == "test" +Requires-Dist: elasticsearch<9,>=8.0.1; extra == "test" +Requires-Dist: xlwt; extra == "test" +Requires-Dist: tablib>=3.2.0; extra == "test" +Requires-Dist: sqliteschema>=2; extra == "test" +Requires-Dist: loguru<1,>=0.4.1; extra == "test" +Requires-Dist: XlsxWriter<4,>=0.9.6; extra == "test" +Requires-Dist: pandas<3,>=0.25.3; extra == "test" +Provides-Extra: theme +Requires-Dist: pytablewriter-altrow-theme<1,>=0.2.0; extra == "theme" +Requires-Dist: pytablewriter-altcol-theme<1,>=0.1.0; extra == "theme" +Provides-Extra: toml +Requires-Dist: toml<1,>=0.9.3; extra == "toml" +Provides-Extra: yaml +Requires-Dist: PyYAML<7,>=3.11; extra == "yaml" + +.. contents:: **pytablewriter** + :backlinks: top + :depth: 2 + +Summary +========= +`pytablewriter `__ is a Python library to write a table in various formats: AsciiDoc / CSV / Elasticsearch / HTML / JavaScript / JSON / LaTeX / LDJSON / LTSV / Markdown / MediaWiki / NumPy / Excel / Pandas / Python / reStructuredText / SQLite / TOML / TSV / YAML. + +.. image:: https://badge.fury.io/py/pytablewriter.svg + :target: https://badge.fury.io/py/pytablewriter + :alt: PyPI package version + +.. image:: https://anaconda.org/conda-forge/pytablewriter/badges/version.svg + :target: https://anaconda.org/conda-forge/pytablewriter + :alt: conda-forge package version + +.. image:: https://img.shields.io/pypi/pyversions/pytablewriter.svg + :target: https://pypi.org/project/pytablewriter/ + :alt: Supported Python versions + +.. image:: https://img.shields.io/pypi/implementation/pytablewriter.svg + :target: https://pypi.org/project/pytablewriter + :alt: Supported Python implementations + +.. image:: https://github.com/thombashi/pytablewriter/actions/workflows/ci.yml/badge.svg + :target: https://github.com/thombashi/pytablewriter/actions/workflows/ci.yml + :alt: CI status of Linux/macOS/Windows + +.. image:: https://coveralls.io/repos/github/thombashi/pytablewriter/badge.svg?branch=master + :target: https://coveralls.io/github/thombashi/pytablewriter?branch=master + :alt: Test coverage + +.. image:: https://github.com/thombashi/pytablewriter/actions/workflows/github-code-scanning/codeql/badge.svg + :target: https://github.com/thombashi/pytablewriter/actions/workflows/github-code-scanning/codeql + :alt: CodeQL + +Features +-------- +- Write a table in various formats: + - Text formats: + - `AsciiDoc `__ + - CSV / Tab-separated values (TSV) / Space-separated values (SSV) + - HTML / CSS + - JSON / `Line-delimited JSON(LDJSON) `__ + - `Labeled Tab-separated Values (LTSV) `__ + - LaTeX: ``tabular``/``array`` environment + - Markdown: CommonMark / `GitHub Flavored Markdown (GFM) `__ / `kramdown `__ + - `MediaWiki `__ + - reStructuredText: `Grid Tables `__/`Simple Tables `__/`CSV Table `__ + - Source code (definition of a variable that represents tabular data) + - JavaScript / `NumPy `__ (`numpy.array `__) / `Pandas `__ (`pandas.DataFrame `__) / Python + - `TOML `__ + - `YAML `__ + - Unicode + - Binary file formats: + - Microsoft Excel :superscript:`TM` (``.xlsx``/``.xls`` file format) + - `pandas.DataFrame `__ pickle file + - `SQLite `__ database + - Application-specific formats: + - `Elasticsearch `__ +- Automatic table cell formatting: + - Alignment + - Padding + - Decimal places of numbers +- Customize table cell styles: + - Text/Background color + - Text alignment + - Font size/weight + - Thousand separator for numbers: e.g. ``1,000``/``1 000`` +- Configure output: + - Write a table to a stream such as a file/standard-output/string-buffer/Jupyter-Notebook + - Get rendered tabular text +- Data sources: + - nested list + - CSV + - `pandas.DataFrame `__ / `pandas.Series `__ + - etc. +- Multibyte character support +- ANSI color support + +Installation +============ + +Installation: pip +------------------------------ +:: + + pip install pytablewriter + +Some of the formats require additional dependency packages, you can install these packages as follows: + +.. csv-table:: Installation of optional dependencies + :header: Installation example, Remark + + ``pip install pytablewriter[es]``, Elasticsearch + ``pip install pytablewriter[excel]``, Excel + ``pip install pytablewriter[html]``, HTML + ``pip install pytablewriter[sqlite]``, SQLite database + ``pip install pytablewriter[toml]``, TOML + ``pip install pytablewriter[theme]``, pytablewriter theme plugins + ``pip install pytablewriter[all]``, Install all of the optional dependencies + +Installation: conda +------------------------------ +:: + + conda install -c conda-forge pytablewriter + +Installation: apt +------------------------------ +:: + + sudo add-apt-repository ppa:thombashi/ppa + sudo apt update + sudo apt install python3-pytablewriter + +Examples +========== +Write tables +-------------- +Write a Markdown table +~~~~~~~~~~~~~~~~~~~~~~~~ +:Sample Code: + .. code-block:: python + + from pytablewriter import MarkdownTableWriter + + def main(): + writer = MarkdownTableWriter( + table_name="example_table", + headers=["int", "float", "str", "bool", "mix", "time"], + value_matrix=[ + [0, 0.1, "hoge", True, 0, "2017-01-01 03:04:05+0900"], + [2, "-2.23", "foo", False, None, "2017-12-23 45:01:23+0900"], + [3, 0, "bar", "true", "inf", "2017-03-03 33:44:55+0900"], + [-10, -9.9, "", "FALSE", "nan", "2017-01-01 00:00:00+0900"], + ], + ) + writer.write_table() + + if __name__ == "__main__": + main() + +:Output: + .. code-block:: + + # example_table + |int|float|str |bool | mix | time | + |--:|----:|----|-----|-------:|------------------------| + | 0| 0.10|hoge|True | 0|2017-01-01 03:04:05+0900| + | 2|-2.23|foo |False| |2017-12-23 12:34:51+0900| + | 3| 0.00|bar |True |Infinity|2017-03-03 22:44:55+0900| + |-10|-9.90| |False| NaN|2017-01-01 00:00:00+0900| + +:Rendering Result: + .. figure:: https://cdn.jsdelivr.net/gh/thombashi/pytablewriter@master/docs/pages/examples/table_format/text/ss/markdown.png + :alt: https://github.com/thombashi/pytablewriter/blob/master/docs/pages/examples/table_format/text/ss/markdown.png + + Rendered markdown at GitHub + +Write a Markdown table with margins +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +:Sample Code: + .. code-block:: python + + from pytablewriter import MarkdownTableWriter + + def main(): + writer = MarkdownTableWriter( + table_name="write a table with margins", + headers=["int", "float", "str", "bool", "mix", "time"], + value_matrix=[ + [0, 0.1, "hoge", True, 0, "2017-01-01 03:04:05+0900"], + [2, "-2.23", "foo", False, None, "2017-12-23 45:01:23+0900"], + [3, 0, "bar", "true", "inf", "2017-03-03 33:44:55+0900"], + [-10, -9.9, "", "FALSE", "nan", "2017-01-01 00:00:00+0900"], + ], + margin=1 # add a whitespace for both sides of each cell + ) + writer.write_table() + + if __name__ == "__main__": + main() + +:Output: + .. code-block:: + + # write a table with margins + | int | float | str | bool | mix | time | + | --: | ----: | ---- | ----- | -------: | ------------------------ | + | 0 | 0.10 | hoge | True | 0 | 2017-01-01 03:04:05+0900 | + | 2 | -2.23 | foo | False | | 2017-12-23 12:34:51+0900 | + | 3 | 0.00 | bar | True | Infinity | 2017-03-03 22:44:55+0900 | + | -10 | -9.90 | | False | NaN | 2017-01-01 00:00:00+0900 | + +``margin`` attribute can be available for all of the text format writer classes. + +Write a GitHub Flavored Markdown (GFM) table +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +If you set ``flavor`` keyword argument of ``MarkdownTableWriter`` class to ``"github"`` or ``"gfm"``, the writer will output markdown tables with GitHub flavor. +GFM can apply some additional styles to tables such as ``fg_color`` (text color). + +:Sample Code: + .. code-block:: python + + from pytablewriter import MarkdownTableWriter + from pytablewriter.style import Style + + writer = MarkdownTableWriter( + column_styles=[ + Style(fg_color="red"), + Style(fg_color="green", decoration_line="underline"), + ], + headers=["A", "B"], + value_matrix=[ + ["abc", 1], + ["efg", 2], + ], + margin=1, + flavor="github", + enable_ansi_escape=False, + ) + writer.write_table() + +Rendered results can be found at `here `__ + +Apply styles to GFM table with programmatically +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +Applying style filters to GFM allows for more flexible style settings for cells. +See also the `example <#style-filter>`_ + +Write a Markdown table to a stream or a file +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +`Refer an example `__ + +Write a table to an Excel sheet +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +:Sample Code: + .. code-block:: python + + from pytablewriter import ExcelXlsxTableWriter + + def main(): + writer = ExcelXlsxTableWriter() + writer.table_name = "example" + writer.headers = ["int", "float", "str", "bool", "mix", "time"] + writer.value_matrix = [ + [0, 0.1, "hoge", True, 0, "2017-01-01 03:04:05+0900"], + [2, "-2.23", "foo", False, None, "2017-12-23 12:34:51+0900"], + [3, 0, "bar", "true", "inf", "2017-03-03 22:44:55+0900"], + [-10, -9.9, "", "FALSE", "nan", "2017-01-01 00:00:00+0900"], + ] + writer.dump("sample.xlsx") + + if __name__ == "__main__": + main() + +:Output: + .. figure:: https://cdn.jsdelivr.net/gh/thombashi/pytablewriter@master/docs/pages/examples/table_format/binary/spreadsheet/ss/excel_single.png + :alt: https://github.com/thombashi/pytablewriter/blob/master/docs/pages/examples/table_format/binary/spreadsheet/ss/excel_single.png + + Output excel file (``sample_single.xlsx``) + +Write a Unicode table +~~~~~~~~~~~~~~~~~~~~~~~ +:Sample Code: + .. code-block:: python + + from pytablewriter import UnicodeTableWriter + + def main(): + writer = UnicodeTableWriter( + table_name="example_table", + headers=["int", "float", "str", "bool", "mix", "time"], + value_matrix=[ + [0, 0.1, "hoge", True, 0, "2017-01-01 03:04:05+0900"], + [2, "-2.23", "foo", False, None, "2017-12-23 45:01:23+0900"], + [3, 0, "bar", "true", "inf", "2017-03-03 33:44:55+0900"], + [-10, -9.9, "", "FALSE", "nan", "2017-01-01 00:00:00+0900"], + ] + ) + writer.write_table() + + if __name__ == "__main__": + main() + +:Output: + .. code-block:: + + ┌───┬─────┬────┬─────┬────────┬────────────────────────┐ + │int│float│str │bool │ mix │ time │ + ├───┼─────┼────┼─────┼────────┼────────────────────────┤ + │ 0│ 0.10│hoge│True │ 0│2017-01-01 03:04:05+0900│ + ├───┼─────┼────┼─────┼────────┼────────────────────────┤ + │ 2│-2.23│foo │False│ │2017-12-23 12:34:51+0900│ + ├───┼─────┼────┼─────┼────────┼────────────────────────┤ + │ 3│ 0.00│bar │True │Infinity│2017-03-03 22:44:55+0900│ + ├───┼─────┼────┼─────┼────────┼────────────────────────┤ + │-10│-9.90│ │False│ NaN│2017-01-01 00:00:00+0900│ + └───┴─────┴────┴─────┴────────┴────────────────────────┘ + +Write a table with JavaScript format (as a nested list variable definition) +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +:Sample Code: + .. code-block:: python + + import pytablewriter as ptw + + + def main(): + writer = ptw.JavaScriptTableWriter( + table_name="js_variable", + headers=["int", "float", "str", "bool", "mix", "time"], + value_matrix=[ + [0, 0.1, "hoge", True, 0, "2017-01-01 03:04:05+0900"], + [2, "-2.23", "foo", False, None, "2017-12-23 45:01:23+0900"], + [3, 0, "bar", "true", "inf", "2017-03-03 33:44:55+0900"], + [-10, -9.9, "", "FALSE", "nan", "2017-01-01 00:00:00+0900"], + ], + ) + + writer.write_table() + + + if __name__ == "__main__": + main() + +:Output: + .. code-block:: js + + const js_variable = [ + ["int", "float", "str", "bool", "mix", "time"], + [0, 0.1, "hoge", true, 0, "2017-01-01 03:04:05+0900"], + [2, -2.23, "foo", false, null, "2017-12-23 45:01:23+0900"], + [3, 0, "bar", true, Infinity, "2017-03-03 33:44:55+0900"], + [-10, -9.9, "", "FALSE", NaN, "2017-01-01 00:00:00+0900"] + ]; + +Write a Markdown table from ``pandas.DataFrame`` instance +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +``from_dataframe`` method of writer classes will set up tabular data from ``pandas.DataFrame``: + +:Sample Code: + .. code-block:: python + + from textwrap import dedent + import pandas as pd + import io + from pytablewriter import MarkdownTableWriter + + def main(): + csv_data = io.StringIO(dedent("""\ + "i","f","c","if","ifc","bool","inf","nan","mix_num","time" + 1,1.10,"aa",1.0,"1",True,Infinity,NaN,1,"2017-01-01 00:00:00+09:00" + 2,2.20,"bbb",2.2,"2.2",False,Infinity,NaN,Infinity,"2017-01-02 03:04:05+09:00" + 3,3.33,"cccc",-3.0,"ccc",True,Infinity,NaN,NaN,"2017-01-01 00:00:00+09:00" + """)) + df = pd.read_csv(csv_data, sep=',') + + writer = MarkdownTableWriter(dataframe=df) + writer.write_table() + + if __name__ == "__main__": + main() + +:Output: + .. code-block:: + + | i | f | c | if |ifc|bool | inf |nan|mix_num | time | + |--:|---:|----|---:|---|-----|--------|---|-------:|-------------------------| + | 1|1.10|aa | 1.0| 1|True |Infinity|NaN| 1|2017-01-01 00:00:00+09:00| + | 2|2.20|bbb | 2.2|2.2|False|Infinity|NaN|Infinity|2017-01-02 03:04:05+09:00| + | 3|3.33|cccc|-3.0|ccc|True |Infinity|NaN| NaN|2017-01-01 00:00:00+09:00| + + +Adding a column of the DataFrame index if you specify ``add_index_column=True``: + +:Sample Code: + .. code-block:: python + + import pandas as pd + import pytablewriter as ptw + + def main(): + writer = ptw.MarkdownTableWriter(table_name="add_index_column") + writer.from_dataframe( + pd.DataFrame({"A": [1, 2], "B": [10, 11]}, index=["a", "b"]), + add_index_column=True, + ) + writer.write_table() + + if __name__ == "__main__": + main() + +:Output: + .. code-block:: + + # add_index_column + | | A | B | + |---|--:|--:| + |a | 1| 10| + |b | 2| 11| + +Write a Markdown table from space-separated values +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +:Sample Code: + .. code-block:: python + + import pytablewriter as ptw + + + def main(): + writer = ptw.MarkdownTableWriter(table_name="ps") + writer.from_csv( + """ + USER PID %CPU %MEM VSZ RSS TTY STAT START TIME COMMAND + root 1 0.0 0.4 77664 8784 ? Ss May11 0:02 /sbin/init + root 2 0.0 0.0 0 0 ? S May11 0:00 [kthreadd] + root 4 0.0 0.0 0 0 ? I< May11 0:00 [kworker/0:0H] + root 6 0.0 0.0 0 0 ? I< May11 0:00 [mm_percpu_wq] + root 7 0.0 0.0 0 0 ? S May11 0:01 [ksoftirqd/0] + """, + delimiter=" ", + ) + writer.write_table() + + + if __name__ == "__main__": + main() + +:Output: + .. code-block:: + + # ps + |USER|PID|%CPU|%MEM| VSZ |RSS |TTY|STAT|START|TIME| COMMAND | + |----|--:|---:|---:|----:|---:|---|----|-----|----|--------------| + |root| 1| 0| 0.4|77664|8784|? |Ss |May11|0:02|/sbin/init | + |root| 2| 0| 0.0| 0| 0|? |S |May11|0:00|[kthreadd] | + |root| 4| 0| 0.0| 0| 0|? |I< |May11|0:00|[kworker/0:0H]| + |root| 6| 0| 0.0| 0| 0|? |I< |May11|0:00|[mm_percpu_wq]| + |root| 7| 0| 0.0| 0| 0|? |S |May11|0:01|[ksoftirqd/0] | + +Get rendered tabular text as str +---------------------------------- +``dumps`` method returns rendered tabular text. +``dumps`` only available for text format writers. + +:Sample Code: + .. code-block:: python + + import pytablewriter as ptw + + + def main(): + writer = ptw.MarkdownTableWriter( + headers=["int", "float", "str", "bool", "mix", "time"], + value_matrix=[ + [0, 0.1, "hoge", True, 0, "2017-01-01 03:04:05+0900"], + [2, "-2.23", "foo", False, None, "2017-12-23 45:01:23+0900"], + [3, 0, "bar", "true", "inf", "2017-03-03 33:44:55+0900"], + [-10, -9.9, "", "FALSE", "nan", "2017-01-01 00:00:00+0900"], + ], + ) + + print(writer.dumps()) + + + if __name__ == "__main__": + main() + +:Output: + .. code-block:: + + |int|float|str |bool | mix | time | + |--:|----:|----|-----|-------:|------------------------| + | 0| 0.10|hoge|True | 0|2017-01-01 03:04:05+0900| + | 2|-2.23|foo |False| |2017-12-23 45:01:23+0900| + | 3| 0.00|bar |True |Infinity|2017-03-03 33:44:55+0900| + |-10|-9.90| |False| NaN|2017-01-01 00:00:00+0900| + +Configure table styles +------------------------ +Column styles +~~~~~~~~~~~~~~~ +Writers can specify +`Style `__ +for each column by ``column_styles`` attribute of writer classes. + +:Sample Code: + .. code-block:: python + + import pytablewriter as ptw + from pytablewriter.style import Style + + + def main(): + writer = ptw.MarkdownTableWriter( + table_name="set style by column_styles", + headers=[ + "auto align", + "left align", + "center align", + "bold", + "italic", + "bold italic ts", + ], + value_matrix=[ + [11, 11, 11, 11, 11, 11], + [1234, 1234, 1234, 1234, 1234, 1234], + ], + column_styles=[ + Style(), + Style(align="left"), + Style(align="center"), + Style(font_weight="bold"), + Style(font_style="italic"), + Style(font_weight="bold", font_style="italic", thousand_separator=","), + ], # specify styles for each column + ) + writer.write_table() + + + if __name__ == "__main__": + main() + +:Output: + .. code-block:: + + # set style by styles + |auto align|left align|center align| bold |italic|bold italic ts| + |---------:|----------|:----------:|-------:|-----:|-------------:| + | 11|11 | 11 | **11**| _11_| _**11**_| + | 1234|1234 | 1234 |**1234**|_1234_| _**1,234**_| + + `Rendering result `__ + + +You can also set ``Style`` to a specific column with an index or header by using ``set_style`` method: + +:Sample Code: + .. code-block:: python + + from pytablewriter import MarkdownTableWriter + from pytablewriter.style import Style + + def main(): + writer = MarkdownTableWriter() + writer.headers = ["A", "B", "C",] + writer.value_matrix = [[11, 11, 11], [1234, 1234, 1234]] + + writer.table_name = "set style by column index" + writer.set_style(1, Style(align="center", font_weight="bold")) + writer.set_style(2, Style(thousand_separator=" ")) + writer.write_table() + writer.write_null_line() + + writer.table_name = "set style by header" + writer.set_style("B", Style(font_style="italic")) + writer.write_table() + + if __name__ == "__main__": + main() + +:Output: + .. code-block:: + + # set style by column index + | A | B | C | + |---:|:------:|----:| + | 11| **11** | 11| + |1234|**1234**|1 234| + + # set style by header + | A | B | C | + |---:|-----:|----:| + | 11| _11_| 11| + |1234|_1234_|1 234| + +Style filter +~~~~~~~~~~~~~~ +You can apply styles to specific cells by using style filters. +Style filters will be written as Python functions. +Examples of a style filter function and how you apply it are as follows: + +:Sample Code: + .. code-block:: python + + from typing import Any, Optional + + from pytablewriter import MarkdownTableWriter + from pytablewriter.style import Cell, Style + + + def style_filter(cell: Cell, **kwargs: Any) -> Optional[Style]: + if cell.is_header_row(): + return None + + if cell.col == 0: + return Style(font_weight="bold") + + value = int(cell.value) + + if value > 80: + return Style(fg_color="red", font_weight="bold", decoration_line="underline") + elif value > 50: + return Style(fg_color="yellow", font_weight="bold") + elif value > 20: + return Style(fg_color="green") + + return Style(fg_color="lightblue") + + + writer = MarkdownTableWriter( + table_name="style filter example", + headers=["Key", "Value 1", "Value 2"], + value_matrix=[ + ["A", 95, 40], + ["B", 55, 5], + ["C", 30, 85], + ["D", 0, 69], + ], + flavor="github", + enable_ansi_escape=False, + ) + writer.add_style_filter(style_filter) + writer.write_table() + +Rendered results can be found at `here `__ + +Theme +~~~~~~~ +`Theme ` +consists of a set of style filters. +The following command will install external predefined themes: + +:: + + pip install pytablewriter[theme] + +Themes can be set via the constructor of the writer classes or the ``set_theme`` method. +The following is an example of setting the ``altrow`` theme via the constructor. +``altrow`` theme will be colored rows alternatively: + +:Sample Code: + .. code-block:: python + + import pytablewriter as ptw + + writer = ptw.TableWriterFactory.create_from_format_name( + "markdown", + headers=["INT", "STR"], + value_matrix=[[1, "hoge"], [2, "foo"], [3, "bar"]], + margin=1, + theme="altrow", + ) + writer.write_table() + +:Output: + .. figure:: https://cdn.jsdelivr.net/gh/thombashi/pytablewriter-altrow-theme@master/ss/ptw-altrow-theme_example_default.png + :alt: https://github.com/thombashi/pytablewriter-altrow-theme/blob/master/ss/ptw-altrow-theme_example_default.png + +`[theme]` extras includes the following themes: + +- `pytablewriter-altrow-theme `__ + - `Generated HTML table example `__ +- `pytablewriter-altcol-theme `__ + - `Generated HTML table example `__ + +Make tables for specific applications +--------------------------------------- +Render a table on Jupyter Notebook +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +All table writer class instances in ``pytablewriter`` can render in Jupyter Notebook. +To render writers at notebook cells, you will require the dependency packages to be installed either by: + +- ``pip install pytablewriter[html]`` or +- ``pip install pytablewriter[all]`` + +Jupyter Notebook code examples can be found `here `__: + +.. figure:: https://cdn.jsdelivr.net/gh/thombashi/pytablewriter@master/ss/jupyter_notebook.png + :alt: https://github.com/thombashi/pytablewriter/blob/master/ss/jupyter_notebook.png + + Table rendering results of Jupyter Notebook + +Multibyte character support +----------------------------- +Write a table using multibyte character +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +You can use multibyte characters as table data. +Multibyte characters are also properly padded and aligned. + +:Sample Code: + .. code-block:: python + + import pytablewriter as ptw + + + def main(): + writer = ptw.RstSimpleTableWriter( + table_name="生成に関するパターン", + headers=["パターン名", "概要", "GoF", "Code Complete[1]"], + value_matrix=[ + ["Abstract Factory", "関連する一連のインスタンスを状況に応じて、適切に生成する方法を提供する。", "Yes", "Yes"], + ["Builder", "複合化されたインスタンスの生成過程を隠蔽する。", "Yes", "No"], + ["Factory Method", "実際に生成されるインスタンスに依存しない、インスタンスの生成方法を提供する。", "Yes", "Yes"], + ["Prototype", "同様のインスタンスを生成するために、原型のインスタンスを複製する。", "Yes", "No"], + ["Singleton", "あるクラスについて、インスタンスが単一であることを保証する。", "Yes", "Yes"], + ], + ) + writer.write_table() + + + if __name__ == "__main__": + main() + +:Output: + .. figure:: https://cdn.jsdelivr.net/gh/thombashi/pytablewriter@master/docs/pages/examples/multibyte/ss/multi_byte_char.png + :alt: https://github.com/thombashi/pytablewriter/blob/master/docs/pages/examples/multibyte/ss/multi_byte_char.png + + Output of multi-byte character table + +Multiprocessing +----------------- +You can increase the number of workers to process table data via ``max_workers`` attribute of a writer. +The more ``max_workers`` the less processing time when tabular data is large and the execution environment has available cores. + +If you increase ``max_workers`` larger than one, recommend using main guarded as follows to avoid problems caused by multi-processing: + +.. code-block:: python + + from multiprocessing import cpu_count + import pytablewriter as ptw + + def main(): + writer = ptw.MarkdownTableWriter() + writer.max_workers = cpu_count() + ... + + if __name__ == "__main__": + main() + +For more information +---------------------- +More examples are available at +https://pytablewriter.rtfd.io/en/latest/pages/examples/index.html + +Dependencies +============ +- Python 3.9+ +- `Python package dependencies (automatically installed) `__ + + +Optional dependencies +--------------------- +- ``logging`` extras + - `loguru `__: Used for logging if the package installed +- ``from`` extras + - `pytablereader `__ +- ``es`` extra + - `elasticsearch `__ +- ``excel`` extras + - `xlwt `__ + - `XlsxWriter `__ +- ``html`` extras + - `dominate `__ +- ``sqlite`` extras + - `SimpleSQLite `__ +- ``theme`` extras + - `pytablewriter-altrow-theme `__ + - `pytablewriter-altcol-theme `__ +- ``toml`` extras + - `toml `__ + +Documentation +=============== +https://pytablewriter.rtfd.io/ + +Projects using pytablewriter +================================== +- `pytest-md-report `__ + + +Related Projects +================================== +- `pytablereader `__ + - Tabular data loaded by ``pytablereader`` can be written another tabular data format with ``pytablewriter``. + +Sponsors +==================================== +|chasbecker| |shiguredo| |b4tman| |Arturi0| |github| + +.. |chasbecker| image:: https://avatars.githubusercontent.com/u/44389260?s=48&u=6da7176e51ae2654bcfd22564772ef8a3bb22318&v=4 + :target: https://github.com/chasbecker + :alt: ex-sponsor: Charles Becker (chasbecker) +.. |shiguredo| image:: https://avatars.githubusercontent.com/u/2549434?s=48&v=4 + :target: https://github.com/shiguredo + :alt: ex-sponsor: 時雨堂 (shiguredo) +.. |b4tman| image:: https://avatars.githubusercontent.com/u/3658062?s=48&v=4 + :target: https://github.com/b4tman + :alt: onetime: Dmitry Belyaev (b4tman) +.. |Arturi0| image:: https://avatars.githubusercontent.com/u/46711571?s=48&u=57687c0e02d5d6e8eeaf9177f7b7af4c9f275eb5&v=4 + :target: https://github.com/Arturi0 + :alt: onetime: Arturi0 +.. |github| image:: https://avatars.githubusercontent.com/u/9919?s=48&v=4 + :target: https://github.com/github + :alt: onetime: GitHub (github) + +`Become a sponsor `__ + diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/pytablewriter-1.2.1.dist-info/RECORD b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/pytablewriter-1.2.1.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..5c67b3693aa9a1d2fa7daebedddbff3a6490c552 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/pytablewriter-1.2.1.dist-info/RECORD @@ -0,0 +1,138 @@ +pytablewriter-1.2.1.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +pytablewriter-1.2.1.dist-info/LICENSE,sha256=Ewo1uRffGVGu-_NmEgSR1RJARah97RI_IN7SFSH046I,1089 +pytablewriter-1.2.1.dist-info/METADATA,sha256=Egd2rPMhFcp60Ff2zC2VEmUjSIV6V6Vnn7mTDeb8PR4,38178 +pytablewriter-1.2.1.dist-info/RECORD,, +pytablewriter-1.2.1.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +pytablewriter-1.2.1.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91 +pytablewriter-1.2.1.dist-info/top_level.txt,sha256=4qovxzrpT62Feu8LLdPGtIqYBswTr4QcU4mRmpM61-k,14 +pytablewriter/__init__.py,sha256=E2Y4TxopUWgqMateYeM22S6pGZct8qa_S78a1J_x9ao,2942 +pytablewriter/__pycache__/__init__.cpython-312.pyc,, +pytablewriter/__pycache__/__version__.cpython-312.pyc,, +pytablewriter/__pycache__/_converter.cpython-312.pyc,, +pytablewriter/__pycache__/_factory.cpython-312.pyc,, +pytablewriter/__pycache__/_function.cpython-312.pyc,, +pytablewriter/__pycache__/_table_format.cpython-312.pyc,, +pytablewriter/__pycache__/error.cpython-312.pyc,, +pytablewriter/__version__.py,sha256=Yl4cZCSg_uar8nSXzhiZUTsEIwUveZ9Wn1et0RO9RL8,268 +pytablewriter/_converter.py,sha256=iPlzCNzbGPJ4eSfgMz7DwD7GjaV0n1zxBm_iIzbvG7E,238 +pytablewriter/_factory.py,sha256=JEzLpd5ri6DqGynTvMQ7DVhK36aoKKe-IedMxXjyVds,10847 +pytablewriter/_function.py,sha256=HkZK5rZRRzzOG0qCtLfoLQE8s6ruifdMvQYXePOUgrw,2396 +pytablewriter/_logger/__init__.py,sha256=DzORajZGSzcVR5wMlNgQ2b54Pr1CBgaN3OycGTp9s7g,107 +pytablewriter/_logger/__pycache__/__init__.cpython-312.pyc,, +pytablewriter/_logger/__pycache__/_logger.cpython-312.pyc,, +pytablewriter/_logger/__pycache__/_null_logger.cpython-312.pyc,, +pytablewriter/_logger/_logger.py,sha256=OFQ-ZwE7Ly_HCW2fATCowChDvkFF2G_CjFwu6A2f5v8,3311 +pytablewriter/_logger/_null_logger.py,sha256=QJuaErUIV_x6NjQ9qNX9eNSi_GB_9CrO7lKeXYZnuaw,1088 +pytablewriter/_table_format.py,sha256=0OMnMrF0H9Wms7FKQM-JaMGob2qjgaIKWb3shQqXeuM,9479 +pytablewriter/error.py,sha256=MwPbc4EtUklc7X3eVWADVCA6rrzelfsBcH16E0pJQeE,787 +pytablewriter/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +pytablewriter/sanitizer/__init__.py,sha256=Ob9cbVV0DBI6W6fupmMIHEgoSCdaGeyxo_VhfvNizEM,703 +pytablewriter/sanitizer/__pycache__/__init__.cpython-312.pyc,, +pytablewriter/sanitizer/__pycache__/_base.cpython-312.pyc,, +pytablewriter/sanitizer/__pycache__/_elasticsearch.cpython-312.pyc,, +pytablewriter/sanitizer/__pycache__/_excel.cpython-312.pyc,, +pytablewriter/sanitizer/__pycache__/_interface.cpython-312.pyc,, +pytablewriter/sanitizer/__pycache__/_javascript.cpython-312.pyc,, +pytablewriter/sanitizer/__pycache__/_python.cpython-312.pyc,, +pytablewriter/sanitizer/_base.py,sha256=jSpUcAA532hUAB9FFmYmh4Syp2hDlWsb_MnTIoCe6oc,2957 +pytablewriter/sanitizer/_elasticsearch.py,sha256=SB8GCeR_yl509Lz2uyblGnpYI56P4DSFaC3byB0BdQI,732 +pytablewriter/sanitizer/_excel.py,sha256=G5hot3llV_qMyziwOJwjZhSR1rTwqNdjRDEO3-29BnE,2510 +pytablewriter/sanitizer/_interface.py,sha256=YRioPeXRAPnDBgZ7G30CnpjivNwfbaxBemjWYi7P6e8,901 +pytablewriter/sanitizer/_javascript.py,sha256=WpX6Pq9JvcxTxK9nlb3IXl-YglE_YHcfp_PDk5jFMRY,3620 +pytablewriter/sanitizer/_python.py,sha256=GLF9rDttSEklwpnMRzAI6VLQVqext6O9_QGWEcgHqDM,3115 +pytablewriter/style/__init__.py,sha256=OmdQIAKEu8o5E9Xu9fN_kQ1SAtCZZPebFEY8QQjGFpQ,1107 +pytablewriter/style/__pycache__/__init__.cpython-312.pyc,, +pytablewriter/style/__pycache__/_cell.cpython-312.pyc,, +pytablewriter/style/__pycache__/_font.cpython-312.pyc,, +pytablewriter/style/__pycache__/_style.cpython-312.pyc,, +pytablewriter/style/__pycache__/_styler.cpython-312.pyc,, +pytablewriter/style/__pycache__/_styler_interface.cpython-312.pyc,, +pytablewriter/style/__pycache__/_theme.cpython-312.pyc,, +pytablewriter/style/_cell.py,sha256=Ggaq9xm2r_oXUv_W2eV1YLZeI-U0AVsTpAJBfj1Dozw,549 +pytablewriter/style/_font.py,sha256=f3e9bKB83JYu7Yow7EYA_6XCJvqyCSMvjrIXC-Uelfc,341 +pytablewriter/style/_style.py,sha256=QJImo6VRWMSZT64yEqpeOyaVuYYbuOCVrlvrvmGvQ0E,12448 +pytablewriter/style/_styler.py,sha256=ceB6An1OIg-cx6lwPvSwRxUriF0EEUwHpmTOrucil9Y,9963 +pytablewriter/style/_styler_interface.py,sha256=rM1OX8rYIQsk9vtPmbrrcTlf4e0_So2XrHT3L4z1bF8,828 +pytablewriter/style/_theme.py,sha256=CpYbRacbIFWoMLxfG4qp2otbtLXLmQIDWyyp2VSSNcc,2836 +pytablewriter/typehint/__init__.py,sha256=FDTB4uiJDm2b0A6IsYtTVO2Z994tb5o3kcXbwkDDKYQ,545 +pytablewriter/typehint/__pycache__/__init__.cpython-312.pyc,, +pytablewriter/writer/__init__.py,sha256=r0ZSklAeSM84jA4xzvTFaXHVe0Il0GjAQ8vk2_mtplQ,1766 +pytablewriter/writer/__pycache__/__init__.cpython-312.pyc,, +pytablewriter/writer/__pycache__/_common.cpython-312.pyc,, +pytablewriter/writer/__pycache__/_elasticsearch.cpython-312.pyc,, +pytablewriter/writer/__pycache__/_interface.cpython-312.pyc,, +pytablewriter/writer/__pycache__/_msgfy.cpython-312.pyc,, +pytablewriter/writer/__pycache__/_null.cpython-312.pyc,, +pytablewriter/writer/__pycache__/_table_writer.cpython-312.pyc,, +pytablewriter/writer/_common.py,sha256=BjKw-NvsyNQw9D8Zrpg8RyjLjgQjc0QiLbp1bQoGROE,221 +pytablewriter/writer/_elasticsearch.py,sha256=OSfj__QgUmz5TKOdlL9iCZ7256iwLEi6ca8t7q0HrFc,6308 +pytablewriter/writer/_interface.py,sha256=SBLhByN6K5EUi3B8hyWZKMIg14P_lLGHA_WjV-sM2jA,2671 +pytablewriter/writer/_msgfy.py,sha256=Qf3VIhuCfstgGOEaYQswrW9R1lgYFknjw33YZJGNyFo,1777 +pytablewriter/writer/_null.py,sha256=YPBm1lc23wCQbVHuYKPPOTtdlZKfZOBIEWpkuBKQEw4,1590 +pytablewriter/writer/_table_writer.py,sha256=86W952psHB-jIzLN9uhCFjsLVKgugn8w-q56bfBAJiA,42116 +pytablewriter/writer/binary/__init__.py,sha256=akvPmDxtQjvKEac2yx9c-96CURTFx0809iPPskpa25c,281 +pytablewriter/writer/binary/__pycache__/__init__.cpython-312.pyc,, +pytablewriter/writer/binary/__pycache__/_excel.cpython-312.pyc,, +pytablewriter/writer/binary/__pycache__/_excel_workbook.cpython-312.pyc,, +pytablewriter/writer/binary/__pycache__/_interface.cpython-312.pyc,, +pytablewriter/writer/binary/__pycache__/_pandas.cpython-312.pyc,, +pytablewriter/writer/binary/__pycache__/_sqlite.cpython-312.pyc,, +pytablewriter/writer/binary/_excel.py,sha256=BzH-PEkUd90F-e5WvFNTmGiUXo5Om35ntw_Js5sM-3E,15543 +pytablewriter/writer/binary/_excel_workbook.py,sha256=ZQyXi3FOms6z5GhcGq6-18sY9Y3GM2GqRmLfQEyKij0,3954 +pytablewriter/writer/binary/_interface.py,sha256=U48pCiVMUgeYSKCINncSN5Sy9OnYQ90LMhC7Ls1C8O0,1487 +pytablewriter/writer/binary/_pandas.py,sha256=7UVtUNYGIIywe7GyOT8uDYLtYezoK9FZzk_NhwHNYQo,2660 +pytablewriter/writer/binary/_sqlite.py,sha256=ZnXqvidGUri1SM-Cxls1NwgVg9riDaPkFnr9iQjGchQ,2982 +pytablewriter/writer/text/__init__.py,sha256=_rk5sczp6H9sag4PXgKIbxSTrgW8HktmlJqN0cXR01M,1384 +pytablewriter/writer/text/__pycache__/__init__.cpython-312.pyc,, +pytablewriter/writer/text/__pycache__/_asciidoc.cpython-312.pyc,, +pytablewriter/writer/text/__pycache__/_borderless.cpython-312.pyc,, +pytablewriter/writer/text/__pycache__/_common.cpython-312.pyc,, +pytablewriter/writer/text/__pycache__/_css.cpython-312.pyc,, +pytablewriter/writer/text/__pycache__/_csv.cpython-312.pyc,, +pytablewriter/writer/text/__pycache__/_html.cpython-312.pyc,, +pytablewriter/writer/text/__pycache__/_interface.cpython-312.pyc,, +pytablewriter/writer/text/__pycache__/_json.cpython-312.pyc,, +pytablewriter/writer/text/__pycache__/_jsonlines.cpython-312.pyc,, +pytablewriter/writer/text/__pycache__/_latex.cpython-312.pyc,, +pytablewriter/writer/text/__pycache__/_ltsv.cpython-312.pyc,, +pytablewriter/writer/text/__pycache__/_markdown.cpython-312.pyc,, +pytablewriter/writer/text/__pycache__/_mediawiki.cpython-312.pyc,, +pytablewriter/writer/text/__pycache__/_rst.cpython-312.pyc,, +pytablewriter/writer/text/__pycache__/_spacealigned.cpython-312.pyc,, +pytablewriter/writer/text/__pycache__/_text_writer.cpython-312.pyc,, +pytablewriter/writer/text/__pycache__/_toml.cpython-312.pyc,, +pytablewriter/writer/text/__pycache__/_tsv.cpython-312.pyc,, +pytablewriter/writer/text/__pycache__/_unicode.cpython-312.pyc,, +pytablewriter/writer/text/__pycache__/_yaml.cpython-312.pyc,, +pytablewriter/writer/text/_asciidoc.py,sha256=E2MO-Z3yVuF3eviwMqNvgkk5_-1R_u-I31bvH85xVHI,4410 +pytablewriter/writer/text/_borderless.py,sha256=4RhWiSppkS2bRIl8osmqkSst-hwDzaAT-GaSyHyHft4,1010 +pytablewriter/writer/text/_common.py,sha256=1YRanAyjyEgo9muaUM3n9pPieKsX0d5Y-_ktI92B_tA,554 +pytablewriter/writer/text/_css.py,sha256=kXaOogifXdFcDgV83D3zd2k1MlfDotv2gTEjgYaJ6cI,5496 +pytablewriter/writer/text/_csv.py,sha256=MBkHpdC7u9L8p9FNws-IpTECzn2-LpbsT9lg1sKEJzg,1484 +pytablewriter/writer/text/_html.py,sha256=lgCJYXuCtWJ9maF4HooylRaYtwYPPrUpHYHAb_VWEAU,6328 +pytablewriter/writer/text/_interface.py,sha256=Qcwjq6w_dz5Lk7Txr42ESnomW0316-LqPBo1HmcRP7I,642 +pytablewriter/writer/text/_json.py,sha256=kAGiwnbFWyLoqbNZL6at4u_DcS1JrUzJV5wOYQZk_J0,5081 +pytablewriter/writer/text/_jsonlines.py,sha256=si6GB85HhqIaTjBYNjg8xD4HhqvYlCHnOTOSl7xY0G4,1296 +pytablewriter/writer/text/_latex.py,sha256=_w_2c2fimZk2TL402lInpWzb3FtQpH6r6KdWbQaGet4,6336 +pytablewriter/writer/text/_ltsv.py,sha256=xsMAMMU2F5UdznagXnQJbz62-nstSiSbjm7vgHlLm_s,1517 +pytablewriter/writer/text/_markdown.py,sha256=ZfBhMN1CgLt5a-2ZxV4QHpiIs3pY-Sh8Fla3w1BXvkA,6201 +pytablewriter/writer/text/_mediawiki.py,sha256=lNGMwtLClFIMYWQeNTZ52wk8SHiT7FO93C580vPqMqU,3312 +pytablewriter/writer/text/_rst.py,sha256=aiE7zHflwU6RnOGQAowbhTF6onBjmPhx6QvQyaDco4Y,6906 +pytablewriter/writer/text/_spacealigned.py,sha256=osMTS0cvNl8qWthlUkB6noAaKGlBUL02MW-YEvMXEgA,897 +pytablewriter/writer/text/_text_writer.py,sha256=u21ijkBHg3KAtOT58yz5gQWeNTy82hAYxSQiO2RxFNE,20556 +pytablewriter/writer/text/_toml.py,sha256=oUQRIiNIXQ47ccGasVohbDGBksMMxDETv0mnbCngVC8,2265 +pytablewriter/writer/text/_tsv.py,sha256=xLXiOznMZ8W8fKa-xfZCNlTl2Q4_HWFTUQlR4__DjuU,467 +pytablewriter/writer/text/_unicode.py,sha256=-2W2O-FaBkPDAJuwBKLEutGS09y7DcasK4Q83K0GXiE,3532 +pytablewriter/writer/text/_yaml.py,sha256=ERDWX60MWjPRHf0Gs3q-m38n49NGf460YRR0ru_RF0s,1968 +pytablewriter/writer/text/sourcecode/__init__.py,sha256=25ju5UpRUV7DBNsusSj4YLzOLY5akmmEW7gKegSVtu4,297 +pytablewriter/writer/text/sourcecode/__pycache__/__init__.cpython-312.pyc,, +pytablewriter/writer/text/sourcecode/__pycache__/_javascript.cpython-312.pyc,, +pytablewriter/writer/text/sourcecode/__pycache__/_numpy.cpython-312.pyc,, +pytablewriter/writer/text/sourcecode/__pycache__/_pandas.cpython-312.pyc,, +pytablewriter/writer/text/sourcecode/__pycache__/_python.cpython-312.pyc,, +pytablewriter/writer/text/sourcecode/__pycache__/_sourcecode.cpython-312.pyc,, +pytablewriter/writer/text/sourcecode/_javascript.py,sha256=MHxFDde91ekZ8J0SfblFa7yO8vE8zKsKRJVY6hEsPok,4725 +pytablewriter/writer/text/sourcecode/_numpy.py,sha256=XGIHnGIpu7r5RPBhaadPXuPSSOWLAd2rss8OiI2huNM,1902 +pytablewriter/writer/text/sourcecode/_pandas.py,sha256=3wgXSSpD7UVNp9vmhmNb-lgZrdpZ66Ol8KEFeXLMl1k,2550 +pytablewriter/writer/text/sourcecode/_python.py,sha256=IQwuSTULsksIqp0tH6lpHGkuOQfC4IxxGrlmG-CGrbI,2535 +pytablewriter/writer/text/sourcecode/_sourcecode.py,sha256=BDr9hGrzmukvZzO48jsEInmeOSigJSBblhtgO5OEC7Y,2239 diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/pytablewriter-1.2.1.dist-info/REQUESTED b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/pytablewriter-1.2.1.dist-info/REQUESTED new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/pytablewriter-1.2.1.dist-info/WHEEL b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/pytablewriter-1.2.1.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..ae527e7d64811439e61b93aa375defb30e06edfe --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/pytablewriter-1.2.1.dist-info/WHEEL @@ -0,0 +1,5 @@ +Wheel-Version: 1.0 +Generator: setuptools (75.6.0) +Root-Is-Purelib: true +Tag: py3-none-any + diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/pytablewriter-1.2.1.dist-info/top_level.txt b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/pytablewriter-1.2.1.dist-info/top_level.txt new file mode 100644 index 0000000000000000000000000000000000000000..ee171b9ce88fa902be1b06281de4ce7ade6e10f9 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/pytablewriter-1.2.1.dist-info/top_level.txt @@ -0,0 +1 @@ +pytablewriter diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/triton/__init__.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/triton/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3d637cf0b25c543046eb6a62ee6d5f11f1f2c1d4 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/triton/__init__.py @@ -0,0 +1,82 @@ +"""isort:skip_file""" +__version__ = '3.6.0' + +# --------------------------------------- +# Note: import order is significant here. + +# submodules +from .runtime import ( + autotune, + Config, + heuristics, + JITFunction, + KernelInterface, + reinterpret, + TensorWrapper, + OutOfResources, + InterpreterError, + MockTensor, +) +from .runtime.jit import constexpr_function, jit +from .runtime._async_compile import AsyncCompileMode, FutureKernel +from .compiler import compile, CompilationError +from .errors import TritonError +from .runtime._allocation import set_allocator + +from . import language +from . import testing +from . import tools + +must_use_result = language.core.must_use_result + +__all__ = [ + "AsyncCompileMode", + "autotune", + "cdiv", + "CompilationError", + "compile", + "Config", + "constexpr_function", + "FutureKernel", + "heuristics", + "InterpreterError", + "jit", + "JITFunction", + "KernelInterface", + "language", + "MockTensor", + "must_use_result", + "next_power_of_2", + "OutOfResources", + "reinterpret", + "runtime", + "set_allocator", + "TensorWrapper", + "TritonError", + "testing", + "tools", +] + +# ------------------------------------- +# misc. utilities that don't fit well +# into any specific module +# ------------------------------------- + + +@constexpr_function +def cdiv(x: int, y: int): + return (x + y - 1) // y + + +@constexpr_function +def next_power_of_2(n: int): + """Return the smallest power of 2 greater than or equal to n""" + n -= 1 + n |= n >> 1 + n |= n >> 2 + n |= n >> 4 + n |= n >> 8 + n |= n >> 16 + n |= n >> 32 + n += 1 + return n diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/triton/_filecheck.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/triton/_filecheck.py new file mode 100644 index 0000000000000000000000000000000000000000..d96eb3cca718dd10eb097314aae3209aed100139 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/triton/_filecheck.py @@ -0,0 +1,96 @@ +import functools +import os +import inspect +import subprocess +import tempfile + +import triton +from triton.compiler import ASTSource, make_backend +from triton.backends.compiler import GPUTarget +from triton.experimental.gluon._runtime import GluonASTSource +from triton.runtime.jit import create_function_from_signature +from triton._C.libtriton import ir + +# ===-----------------------------------------------------------------------===# +# filecheck_test +# ===-----------------------------------------------------------------------===# + +# Stub target for testing the frontend. +stub_target = GPUTarget("cuda", 100, 32) + +triton_dir = os.path.dirname(__file__) +filecheck_path = os.path.join(triton_dir, "FileCheck") + + +class MatchError(ValueError): + + def __init__(self, message, module_str): + super().__init__(message) + self.module_str = module_str + + def __str__(self): + return f"{super().__str__()}\n{self.module_str}" + + +def run_filecheck(name, module_str, check_template): + with tempfile.TemporaryDirectory() as tempdir: + temp_module = os.path.join(tempdir, "module") + with open(temp_module, "w") as temp: + temp.write(module_str) + + temp_expected = os.path.join(tempdir, "expected") + with open(temp_expected, "w") as temp: + temp.write(check_template) + + try: + subprocess.check_output( + [filecheck_path, temp_expected, "--input-file", temp_module, "--dump-input-context=50"], + stderr=subprocess.STDOUT) + except subprocess.CalledProcessError as error: + decoded = error.output.decode('unicode_escape') + raise ValueError(decoded) + + +def run_parser(kernel_fn, args=(), kwargs={}, target=stub_target): + if "sanitize_overflow" not in kwargs: + kwargs = dict(kwargs) + kwargs["sanitize_overflow"] = False + backend = make_backend(target) + binder = create_function_from_signature( + kernel_fn.signature, + kernel_fn.params, + backend, + ) + + bound_args, specialization, options = binder(*args, **kwargs) + options, signature, constexprs, attrs = kernel_fn._pack_args(backend, kwargs, bound_args, specialization, options) + source_cls = GluonASTSource if kernel_fn.is_gluon() else ASTSource + src = source_cls(kernel_fn, signature, constexprs, attrs) + + context = ir.context() + ir.load_dialects(context) + backend.load_dialects(context) + + codegen_fns = backend.get_codegen_implementation(options) + module_map = backend.get_module_map() + module = src.make_ir(target, options, codegen_fns, module_map, context) + return module + + +def run_filecheck_test(kernel_fn): + assert isinstance(kernel_fn, triton.runtime.JITFunction) + check_template = inspect.getsource(kernel_fn.fn) + if check_template is None: + raise ValueError("kernel function must have a docstring with FileCheck template") + mlir_module = run_parser(kernel_fn) + + run_filecheck("placeholder", mlir_module.str_nodebug(), check_template) + + +def filecheck_test(fn): + + @functools.wraps(fn) + def test_fn(): + run_filecheck_test(fn) + + return test_fn diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/triton/_internal_testing.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/triton/_internal_testing.py new file mode 100644 index 0000000000000000000000000000000000000000..562e8df721ea7e025101429fb5ca25c9702ebf92 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/triton/_internal_testing.py @@ -0,0 +1,272 @@ +import os +import re +import numpy as np +import torch +import triton +import triton.language as tl +from triton import knobs +from typing import Optional, Set, Union +import pytest + +from numpy.random import RandomState +from triton.runtime.jit import TensorWrapper, reinterpret, type_canonicalisation_dict + +int_dtypes = ['int8', 'int16', 'int32', 'int64'] +uint_dtypes = ['uint8', 'uint16', 'uint32', 'uint64'] +integral_dtypes = int_dtypes + uint_dtypes +float_dtypes = ['float16', 'float32', 'float64'] +float_dtypes_with_bfloat16 = float_dtypes + ['bfloat16'] +dtypes = integral_dtypes + float_dtypes +dtypes_with_bfloat16 = dtypes + ['bfloat16'] +torch_float8_dtypes = ['float8_e4m3fn', 'float8_e5m2'] +torch_dtypes = ['bool'] + int_dtypes + ['uint8'] + float_dtypes + ['bfloat16'] +tma_dtypes = sorted(set(dtypes_with_bfloat16) - {"int64", "uint64", "float64"}) + + +def is_interpreter(): + return os.environ.get('TRITON_INTERPRET', '0') == '1' + + +def get_current_target(): + if is_interpreter(): + return None + return triton.runtime.driver.active.get_current_target() + + +def is_cuda(): + target = get_current_target() + return False if target is None else target.backend == "cuda" + + +def is_ampere_or_newer(): + return is_cuda() and torch.cuda.get_device_capability()[0] >= 8 + + +def is_blackwell(): + return is_cuda() and torch.cuda.get_device_capability()[0] == 10 + + +def is_hopper_or_newer(): + return is_cuda() and torch.cuda.get_device_capability()[0] >= 9 + + +def is_hopper(): + return is_cuda() and torch.cuda.get_device_capability()[0] == 9 + + +def is_sm12x(): + return is_cuda() and torch.cuda.get_device_capability()[0] == 12 + + +def is_hip(): + target = get_current_target() + return False if target is None else target.backend == "hip" + + +def is_hip_cdna2(): + target = get_current_target() + return target is not None and target.backend == 'hip' and target.arch == 'gfx90a' + + +def is_hip_cdna3(): + target = get_current_target() + return target is not None and target.backend == 'hip' and target.arch == 'gfx942' + + +def is_hip_cdna4(): + target = get_current_target() + return target is not None and target.backend == 'hip' and target.arch == 'gfx950' + + +def is_hip_gfx11(): + target = get_current_target() + return target is not None and target.backend == 'hip' and 'gfx11' in target.arch + + +def is_hip_gfx12(): + target = get_current_target() + return target is not None and target.backend == 'hip' and 'gfx12' in target.arch + + +def is_hip_gfx1250(): + target = get_current_target() + return target is not None and target.backend == 'hip' and 'gfx1250' in target.arch + + +def is_hip_cdna(): + return is_hip_cdna2() or is_hip_cdna3() or is_hip_cdna4() + + +def get_hip_lds_size(): + return 163840 if is_hip_cdna4() else 65536 + + +def is_xpu(): + target = get_current_target() + return False if target is None else target.backend == "xpu" + + +def get_arch(): + target = get_current_target() + return "" if target is None else str(target.arch) + + +def numpy_random(shape, dtype_str, rs: Optional[RandomState] = None, low=None, high=None): + """ + Override `rs` if you're calling this function twice and don't want the same + result for both calls. + """ + if isinstance(shape, int): + shape = (shape, ) + if rs is None: + rs = RandomState(seed=17) + if dtype_str in int_dtypes + uint_dtypes: + iinfo = np.iinfo(getattr(np, dtype_str)) + low = iinfo.min if low is None else max(low, iinfo.min) + high = iinfo.max if high is None else min(high, iinfo.max) + dtype = getattr(np, dtype_str) + x = rs.randint(low, high, shape, dtype=dtype) + x[x == 0] = 1 # Workaround. Never return zero so tests of division don't error out. + return x + elif dtype_str and 'float8' in dtype_str: + x = rs.randint(20, 40, shape, dtype=np.int8) + return x + elif dtype_str in float_dtypes: + return rs.normal(0, 1, shape).astype(dtype_str) + elif dtype_str == 'bfloat16': + return (rs.normal(0, 1, shape).astype('float32').view('uint32') & np.uint32(0xffff0000)).view('float32') + elif dtype_str in ['bool', 'int1', 'bool_']: + return rs.normal(0, 1, shape) > 0.0 + else: + raise RuntimeError(f'Unknown dtype {dtype_str}') + + +def to_triton(x: np.ndarray, device, dst_type=None) -> Union[TensorWrapper, torch.Tensor]: + ''' + Note: We need dst_type because the type of x can be different from dst_type. + For example: x is of type `float32`, dst_type is `bfloat16`. + If dst_type is None, we infer dst_type from x. + ''' + t = x.dtype.name + if t in uint_dtypes: + signed_type_name = t.lstrip('u') # e.g. "uint16" -> "int16" + x_signed = x.astype(getattr(np, signed_type_name)) + return reinterpret(torch.tensor(x_signed, device=device), getattr(tl, t)) + else: + if dst_type and 'float8' in dst_type: + return reinterpret(torch.tensor(x, device=device), getattr(tl, dst_type)) + if t == 'float32' and dst_type == 'bfloat16': + return torch.tensor(x, device=device).bfloat16() + return torch.tensor(x, device=device) + + +def str_to_triton_dtype(x: str) -> tl.dtype: + return tl.str_to_ty(type_canonicalisation_dict[x], None) + + +def torch_dtype_name(dtype) -> str: + if isinstance(dtype, triton.language.dtype): + return dtype.name + elif isinstance(dtype, torch.dtype): + # 'torch.int64' -> 'int64' + m = re.match(r'^torch\.(\w+)$', str(dtype)) + return m.group(1) + else: + raise TypeError(f'not a triton or torch dtype: {type(dtype)}') + + +def to_numpy(x): + if isinstance(x, TensorWrapper): + return x.base.cpu().numpy().astype(getattr(np, torch_dtype_name(x.dtype))) + elif isinstance(x, torch.Tensor): + if x.dtype is torch.bfloat16: + return x.cpu().float().numpy() + return x.cpu().numpy() + else: + raise ValueError(f"Not a triton-compatible tensor: {x}") + + +def supports_tma(byval_only=False): + if is_interpreter(): + return True + if not is_cuda(): + return False + cuda_version = knobs.nvidia.ptxas.version + min_cuda_version = (12, 0) if byval_only else (12, 3) + cuda_version_tuple = tuple(map(int, cuda_version.split("."))) + assert len(cuda_version_tuple) == 2, cuda_version_tuple + return torch.cuda.get_device_capability()[0] >= 9 and cuda_version_tuple >= min_cuda_version + + +def supports_ws(): + if is_interpreter(): + return True + if not is_cuda(): + return False + return torch.cuda.get_device_capability()[0] >= 9 + + +def tma_skip_msg(byval_only=False): + if byval_only: + return "Requires __grid_constant__ TMA support (NVIDIA Hopper or higher, CUDA 12.0 or higher)" + else: + return "Requires advanced TMA support (NVIDIA Hopper or higher, CUDA 12.3 or higher)" + + +requires_tma = pytest.mark.skipif(not supports_tma(), reason=tma_skip_msg()) + + +def default_alloc_fn(size: int, align: int, _): + return torch.empty(size, dtype=torch.int8, device="cuda") + + +def unwrap_tensor(t: Union[torch.Tensor, triton.runtime.jit.TensorWrapper]) -> torch.Tensor: + if isinstance(t, triton.runtime.jit.TensorWrapper): + return t.base + return t + + +def _fresh_knobs_impl(skipped_attr: Optional[Set[str]] = None): + from triton import knobs + + if skipped_attr is None: + skipped_attr = set() + + monkeypatch = pytest.MonkeyPatch() + + knobs_map = { + name: knobset + for name, knobset in knobs.__dict__.items() + if isinstance(knobset, knobs.base_knobs) and knobset != knobs.base_knobs and name not in skipped_attr + } + + # We store which variables we need to unset below in finally because + # monkeypatch doesn't appear to reset variables that were never set + # before the monkeypatch.delenv call below. + env_to_unset = [] + prev_propagate_env = knobs.propagate_env + + def fresh_function(): + nonlocal env_to_unset + for name, knobset in knobs_map.items(): + setattr(knobs, name, knobset.copy().reset()) + for knob in knobset.knob_descriptors.values(): + if knob.key in os.environ: + monkeypatch.delenv(knob.key, raising=False) + else: + env_to_unset.append(knob.key) + knobs.propagate_env = True + return knobs + + def reset_function(): + for name, knobset in knobs_map.items(): + setattr(knobs, name, knobset) + # `undo` should be placed before `del os.environ` + # Otherwise, it may restore environment variables that monkeypatch deleted + monkeypatch.undo() + for k in env_to_unset: + if k in os.environ: + del os.environ[k] + knobs.propagate_env = prev_propagate_env + + return fresh_function, reset_function diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/triton/_utils.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/triton/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a3cf0ff0cbdb6047e9379cbc38cc44e5c2b52a62 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/triton/_utils.py @@ -0,0 +1,134 @@ +from __future__ import annotations + +from functools import reduce +from typing import Any, Callable, TYPE_CHECKING, Union, List, Dict + +if TYPE_CHECKING: + from .language import core + IterableType = Union[list[Any], tuple[Any, ...], core.tuple, core.tuple_type] + ObjPath = tuple[int, ...] + +TRITON_MAX_TENSOR_NUMEL = 1048576 + + +def get_iterable_path(iterable: IterableType, path: ObjPath) -> Any: + return reduce(lambda a, idx: a[idx], path, iterable) # type: ignore[index] + + +def set_iterable_path(iterable: IterableType, path: tuple[int, ...], val: Any): + from .language import core + assert len(path) != 0 + prev = iterable if len(path) == 1 else get_iterable_path(iterable, path[:-1]) + assert isinstance(prev, core.tuple) + prev._setitem(path[-1], val) + + +def find_paths_if(iterable: Union[IterableType, Any], pred: Callable[[ObjPath, Any], bool]) -> list[ObjPath]: + from .language import core + is_iterable: Callable[[Any], bool] = lambda x: isinstance(x, (list, tuple, core.tuple, core.tuple_type)) + # We need to use dict so that ordering is maintained, while set doesn't guarantee order + ret: dict[ObjPath, None] = {} + + def _impl(path: tuple[int, ...], current: Any): + if is_iterable(current): + for idx, item in enumerate(current): + _impl((*path, idx), item) + elif pred(path, current): + ret[path] = None + + _impl((), iterable) + + return list(ret.keys()) + + +def is_power_of_two(x): + return (x & (x - 1)) == 0 + + +def validate_block_shape(shape: List[int]): + numel = 1 + for i, d in enumerate(shape): + if not isinstance(d, int): + raise TypeError(f"Shape element {i} must have type `constexpr[int]`, got `constexpr[{type(d)}]") + if not is_power_of_two(d): + raise ValueError(f"Shape element {i} must be a power of 2") + numel *= d + + if numel > TRITON_MAX_TENSOR_NUMEL: + raise ValueError(f"numel ({numel}) exceeds triton maximum tensor numel ({TRITON_MAX_TENSOR_NUMEL})") + return numel + + +type_canonicalisation_dict = { + # we canonicalise all bools to be unsigned: + "bool": "u1", + "int1": "u1", + "uint1": "u1", + "i1": "u1", + # floating-point dtypes: + "float8e4nv": "fp8e4nv", + "float8e5": "fp8e5", + "float8e4b15": "fp8e4b15", + "float8_e4m3fn": "fp8e4nv", + "float8e4b8": "fp8e4b8", + "float8_e4m3fnuz": "fp8e4b8", + "float8_e5m2": "fp8e5", + "float8e5b16": "fp8e5b16", + "float8_e5m2fnuz": "fp8e5b16", + "half": "fp16", + "float16": "fp16", + "bfloat16": "bf16", + "float": "fp32", + "float32": "fp32", + "double": "fp64", + "float64": "fp64", + # signed integers: + "int8": "i8", + "int16": "i16", + "int": "i32", + "int32": "i32", + "int64": "i64", + # unsigned integers: + "uint8": "u8", + "uint16": "u16", + "uint32": "u32", + "uint64": "u64", + "void": "void", +} + +for v in list(type_canonicalisation_dict.values()): + type_canonicalisation_dict[v] = v + + +def canonicalize_dtype(dtype): + dtype_str = str(dtype).split(".")[-1] + return type_canonicalisation_dict[dtype_str] + + +def canonicalize_ptr_dtype(dtype, is_const): + return f"{'*k' if is_const else '*'}{canonicalize_dtype(dtype)}" + + +BITWIDTH_DICT: Dict[str, int] = { + **{f"u{n}": n + for n in (1, 8, 16, 32, 64)}, + **{f"i{n}": n + for n in (1, 8, 16, 32, 64)}, + **{f"fp{n}": n + for n in (16, 32, 64)}, + **{f"fp8{suffix}": 8 + for suffix in ("e4nv", "e4b15", "e4b8", "e5", "e5b16")}, + "bf16": 16, + "void": 0, +} + +for k, v in type_canonicalisation_dict.items(): + BITWIDTH_DICT[k] = BITWIDTH_DICT[v] + + +def get_primitive_bitwidth(dtype: str) -> int: + return BITWIDTH_DICT[dtype] + + +def is_namedtuple(val): + return isinstance(val, type) and issubclass(val, tuple) and hasattr(val, "_fields") diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/triton/errors.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/triton/errors.py new file mode 100644 index 0000000000000000000000000000000000000000..3a0a863553b9d62ccb60d0090ee6d66f8b9e3b79 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/triton/errors.py @@ -0,0 +1,5 @@ +"""Base class for all errors raised by Triton""" + + +class TritonError(Exception): + ... diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/triton/knobs.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/triton/knobs.py new file mode 100644 index 0000000000000000000000000000000000000000..b55bd3de3218a55c4682d60ad95a851db1d3eea2 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/triton/knobs.py @@ -0,0 +1,543 @@ +from __future__ import annotations + +import functools +import importlib +import os +import re +import subprocess +import sysconfig +import pathlib + +from dataclasses import dataclass +from contextlib import contextmanager +from typing import cast, Any, Callable, Generator, Generic, Optional, Protocol, Type, TypeVar, TypedDict, TYPE_CHECKING, Union + +from triton._C.libtriton import getenv, getenv_bool # type: ignore + +if TYPE_CHECKING: + from .runtime.cache import CacheManager, RemoteCacheBackend + from .runtime.jit import JitFunctionInfo, KernelParam + from .compiler.compiler import ASTSource, LazyDict, IRSource + + +class Env: + pass + + +env = Env() + +propagate_env: bool = True + + +def setenv(key: str, value: Optional[str]) -> None: + if not propagate_env: + return + + if value is not None: + os.environ[key] = value + elif key in os.environ: + del os.environ[key] + + +def toenv(val: Any) -> Union[None, tuple[Optional[str]]]: + if val is None: + return (None, ) + + t = type(val) + if t is bool: + return ("1" if val else "0", ) + + if t is str: + return (val, ) + + if t is int: + return (str(val), ) + + return None + + +# There's an asymmetry here so that e.g. env_nvidia_tool can be specified with a +# a string but return an NvidiaTool. +SetType = TypeVar("SetType") +GetType = TypeVar("GetType") + +_NOTHING = object() + + +class env_base(Generic[SetType, GetType]): + + def __init__(self, key: str) -> None: + self.key = key + + def __set_name__(self, objclass: Type[object], name: str) -> None: + self.name = name + + def __get__(self, obj: Optional[object], objclass: Optional[Type[object]]) -> GetType: + py_val = obj.__dict__.get(self.name, _NOTHING) + if py_val is _NOTHING: + return self.get() + return self.transform(py_val) + + def get(self) -> GetType: + raise NotImplementedError() + + def __set__(self, obj: object, value: Union[SetType, Env]) -> None: + if isinstance(value, Env): + obj.__dict__.pop(self.name, None) + else: + obj.__dict__[self.name] = value + if env_val := toenv(value): + setenv(self.key, env_val[0]) + + def __delete__(self, obj: object) -> None: + obj.__dict__.pop(self.name, None) + + def transform(self, val: SetType) -> GetType: + # See comment about GetType/SetType in their definition above. Only needed + # if GetType != SetType. + return cast(GetType, val) + + +class env_str(env_base[str, str]): + + def __init__(self, key: str, default: str): + super().__init__(key) + self.default = default + + def get(self) -> str: + return getenv(self.key, self.default) + + +class env_str_callable_default(env_base[str, str]): + + def __init__(self, key: str, default_factory: Callable[[], str]): + super().__init__(key) + self.default_factory = default_factory + + def get(self) -> str: + env_val = getenv(self.key) + if env_val is None: + return self.default_factory() + return env_val + + +class env_bool(env_base[bool, bool]): + + def __init__(self, key: str, default: bool = False) -> None: + super().__init__(key) + self.default = default + + def get(self) -> bool: + return getenv_bool(self.key, self.default) + + +class env_int(env_base[int, int]): + + def __init__(self, key: str, default: int = 0) -> None: + super().__init__(key) + self.default = default + + def get(self) -> int: + val = getenv(self.key) + if val is None: + return self.default + try: + return int(val) + except ValueError as exc: + raise RuntimeError(f"Unable to use {self.key}={val}: expected int") from exc + + +ClassType = TypeVar("ClassType") + + +class env_class(Generic[ClassType], env_base[Optional[Type[ClassType]], Optional[Type[ClassType]]]): + + def __init__(self, key: str, type: str) -> None: + super().__init__(key) + # We can't pass the type directly to avoid import cycles + self.type = type + + def get(self) -> Optional[Type[ClassType]]: + val = getenv(self.key) + if val is None: + return None + comps = val.split(":", 1) + if len(comps) != 2: + raise RuntimeError(f"Unable to read {self.key}: '{val}' isn't of the form MODULE:CLASS") + cls = getattr(importlib.import_module(comps[0]), comps[1]) + + if not any((c.__name__ == self.type for c in cls.mro())): + raise RuntimeError(f"Unable to use '{val}' from {self.key}: not of type '{self.type}'") + + return cast(Type[ClassType], cls) + + +@dataclass +class NvidiaTool: + path: str + version: str + + @staticmethod + @functools.lru_cache + def from_path(path: str) -> Optional[NvidiaTool]: + try: + result = subprocess.check_output([path, "--version"], stderr=subprocess.STDOUT) + version = re.search(r".*release (\d+\.\d+).*", result.decode("utf-8"), flags=re.MULTILINE) + if version is None: + return None + return NvidiaTool(path, version.group(1)) + except (subprocess.CalledProcessError, FileNotFoundError): + return None + + +class env_nvidia_tool(env_base[str, NvidiaTool]): + + def __init__(self, binary: str) -> None: + binary += sysconfig.get_config_var("EXE") + self.binary = binary + self.default_path = os.path.join(os.path.dirname(__file__), "backends", "nvidia", "bin", binary) + # Convert ptxas-blackwell to PTXAS_BLACKWELL, not PTXAS-BLACKWELL + super().__init__(f"TRITON_{binary.upper().replace('-', '_')}_PATH") + + def get(self) -> NvidiaTool: + return self.transform(getenv(self.key)) + + def transform(self, path: str) -> NvidiaTool: + # We still add default as fallback in case the pointed binary isn't + # accessible. + if path is not None: + paths = [path, self.default_path] + else: + paths = [self.default_path] + + for path in paths: + if tool := NvidiaTool.from_path(path): + return tool + + raise RuntimeError(f"Cannot find {self.binary}") + + +# Separate classes so that types are correct +class env_opt_str(env_base[Optional[str], Optional[str]]): + + def get(self) -> Optional[str]: + return getenv(self.key) + + +class env_opt_bool(env_base): + + def get(self) -> Optional[str]: + return getenv_bool(self.key, None) + + +@dataclass(frozen=True) +class CompileTimes: + """ + Model holding timing information for an invocation of the compiler. + + All times in microseconds. + """ + + # Duration of make_ir + ir_initialization: int + + # Ordered mapping from lowering stage to duration spent in that stage. + # Keyed by stage extension, e.g. ttir, ttgir + lowering_stages: list[tuple[str, int]] + + # Duration of saving artifacts/metadata to cache + store_results: int + + @property + def total_lowering(self) -> int: + return sum((stage[1] for stage in self.lowering_stages)) + + @property + def total(self) -> int: + return self.ir_initialization + self.total_lowering + self.store_results + + +class CompilationListener(Protocol): + + def __call__(self, *, src: Union[ASTSource, IRSource], metadata: dict[str, Any], metadata_group: dict[str, str], + times: CompileTimes, cache_hit: bool) -> None: + ... + + +knobs_type = TypeVar("knobs_type", bound='base_knobs') + + +class base_knobs: + + @property + def knob_descriptors(self) -> dict[str, env_base]: + return { + k: v + # data descriptors live on the class object + for k, v in type(self).__dict__.items() + if isinstance(v, env_base) + } + + @property + def knobs(self) -> dict[str, Any]: + return {k: getattr(self, k) for k in self.knob_descriptors.keys()} + + def copy(self: knobs_type) -> knobs_type: + res = type(self)() + res.__dict__.update(self.__dict__) + return res + + def reset(self: knobs_type) -> knobs_type: + for knob in self.knob_descriptors.keys(): + delattr(self, knob) + return self + + @contextmanager + def scope(self) -> Generator[None, None, None]: + try: + initial_env = {knob.key: getenv(knob.key) for knob in self.knob_descriptors.values()} + orig = dict(self.__dict__) + yield + finally: + self.__dict__.clear() + self.__dict__.update(orig) + + for k, v in initial_env.items(): + if v is not None: + os.environ[k] = v + elif k in os.environ: + del os.environ[k] + + +class BuildImpl(Protocol): + + def __call__(self, name: str, src: str, srcdir: str, library_dirs: list[str], include_dirs: list[str], + libraries: list[str], /) -> str: + ... + + +class build_knobs(base_knobs): + """Configuration controlling how the native compiler is invoked""" + cc: env_opt_str = env_opt_str("CC") + + cudacrt_path: env_opt_str = env_opt_str("TRITON_CUDACRT_PATH") + cudart_path: env_opt_str = env_opt_str("TRITON_CUDART_PATH") + + impl: Optional[BuildImpl] = None + + @property + def backend_dirs(self) -> set[str]: + return {path for path in (self.cudacrt_path, self.cudart_path) if path is not None} + + +class redis_knobs(base_knobs): + key_format: env_str = env_str("TRITON_REDIS_KEY_FORMAT", "triton:{key}:{filename}") + host: env_str = env_str("TRITON_REDIS_HOST", "localhost") + port: env_int = env_int("TRITON_REDIS_PORT", 6379) + + +cache: cache_knobs + + +class cache_knobs(base_knobs): + home_dir: env_str = env_str("TRITON_HOME", os.path.expanduser("~/")) + + dump_dir = env_str_callable_default("TRITON_DUMP_DIR", lambda: cache.get_triton_dir("dump")) + override_dir = env_str_callable_default("TRITON_OVERRIDE_DIR", lambda: cache.get_triton_dir("override")) + dir = env_str_callable_default("TRITON_CACHE_DIR", lambda: cache.get_triton_dir("cache")) + + manager_class: env_class[CacheManager] = env_class("TRITON_CACHE_MANAGER", "CacheManager") + remote_manager_class: env_class[RemoteCacheBackend] = env_class("TRITON_REMOTE_CACHE_BACKEND", "RemoteCacheBackend") + + def get_triton_dir(self, dirname: str) -> str: + return os.path.join(self.home_dir, ".triton", dirname) + + +class compilation_knobs(base_knobs): + override: env_bool = env_bool("TRITON_KERNEL_OVERRIDE") + dump_ir: env_bool = env_bool("TRITON_KERNEL_DUMP") + dump_ir_extract_di_local_variables: env_bool = env_bool("LLVM_EXTRACT_DI_LOCAL_VARIABLES") + store_binary_only: env_bool = env_bool("TRITON_STORE_BINARY_ONLY") + always_compile: env_bool = env_bool("TRITON_ALWAYS_COMPILE") + # TODO: Use enum to constrain / 'typecheck' the values + use_ir_loc: env_opt_str = env_opt_str("USE_IR_LOC") + enable_asan: env_bool = env_bool("TRITON_ENABLE_ASAN") + disable_line_info: env_bool = env_bool("TRITON_DISABLE_LINE_INFO") + front_end_debugging: env_bool = env_bool("TRITON_FRONT_END_DEBUGGING") + allow_non_constexpr_globals: env_bool = env_bool("TRITON_ALLOW_NON_CONSTEXPR_GLOBALS") + # Instrumentation mode is checked on every run, which is expensive. + # We cache the value here to avoid the expensive check on every run. + instrumentation_mode: str = env_str("TRITON_INSTRUMENTATION_MODE", "").get() + listener: Union[CompilationListener, None] = None + + +class autotuning_knobs(base_knobs): + cache: env_bool = env_bool("TRITON_CACHE_AUTOTUNING") + print: env_bool = env_bool("TRITON_PRINT_AUTOTUNING") + + +class LaunchHook(Protocol): + """Hook invoked before and after kernel launching + """ + + def __call__(self, metadata: LazyDict) -> None: + ... + + +class InitHandleHook(Protocol): + """Hook invoked around kernel binary/module loading. + module/function can be None for the *start* hook (before loading). + """ + + def __call__( + self, + module: Optional[object], + function: Optional[Callable], + name: str, + metadata_group: dict[str, str], + hash: str, + ) -> None: + ... + + +F = TypeVar("F", bound=Callable) + + +class HookChain(Generic[F]): + """A chain of hooks of the same type F to be called in order. + """ + + def __init__(self, reversed: bool = False): + self.calls: list[F] = [] + self.reversed = reversed + + def add(self, func: F) -> None: + if func not in self.calls: + self.calls.append(func) + + def remove(self, func: F) -> None: + if func in self.calls: + self.calls.remove(func) + + def __call__(self, *args, **kwargs): + for call in self.calls if not self.reversed else reversed(self.calls): + call(*args, **kwargs) + + +# This is of the form [attr_name, attr_val] +# TODO: Use tuple instead of list for better typing. +KernelAttr = list[Union[str, int]] + + +class JITHookCompileInfo(TypedDict): + key: str + signature: dict[KernelParam, str] + device: int + constants: None + num_warps: int + num_ctas: int + num_stages: int + enable_fp_fusion: bool + launch_cooperative_grid: bool + extern_libs: tuple[tuple[str, str], ...] + configs: list[dict[tuple[int, ...], list[KernelAttr]]] + specialization_data: str + is_warmup: bool + + +class JITHook(Protocol): + + def __call__(self, *, key: str, repr: str, fn: JitFunctionInfo, compile: JITHookCompileInfo, is_manual_warmup: bool, + already_compiled: bool) -> Optional[bool]: + ... + + +class PipelineStagesHook(Protocol): + + def __call__(self, stages, options, language, capability): + ... + + +class runtime_knobs(base_knobs): + interpret: env_bool = env_bool("TRITON_INTERPRET") + # debug is on critical path for kernel launches + # avoid repeated reads from env-var by calling get directly + debug: bool = env_bool("TRITON_DEBUG").get() + override_arch: env_opt_str = env_opt_str("TRITON_OVERRIDE_ARCH") + + launch_enter_hook: HookChain[LaunchHook] = HookChain() + launch_exit_hook: HookChain[LaunchHook] = HookChain(reversed=True) + kernel_load_start_hook: HookChain[InitHandleHook] = HookChain() + kernel_load_end_hook: HookChain[InitHandleHook] = HookChain(reversed=True) + + # Hook for inspecting compiled functions and modules + jit_cache_hook: Optional[JITHook] = None + # Hook to signal that a kernel is done compiling and inspect compiled function. + # jit_cache_hook will always be called before compilation and jit_post_compile_hook after. + jit_post_compile_hook: Optional[JITHook] = None + + # Hook for inspecting compiler pipeline stages + add_stages_inspection_hook: Optional[PipelineStagesHook] = None + + +class language_knobs(base_knobs): + fp32_default: env_opt_str = env_opt_str("TRITON_F32_DEFAULT") + default_fp_fusion: env_bool = env_bool("TRITON_DEFAULT_FP_FUSION", True) + + +class nvidia_knobs(base_knobs): + cuobjdump: env_nvidia_tool = env_nvidia_tool("cuobjdump") + nvdisasm: env_nvidia_tool = env_nvidia_tool("nvdisasm") + ptxas: env_nvidia_tool = env_nvidia_tool("ptxas") + ptxas_blackwell: env_nvidia_tool = env_nvidia_tool("ptxas-blackwell") + + dump_nvptx: env_bool = env_bool("NVPTX_ENABLE_DUMP") + disable_ptxas_opt: env_bool = env_bool("DISABLE_PTXAS_OPT") + ptxas_options: env_opt_str = env_opt_str("PTXAS_OPTIONS") + mock_ptx_version: env_opt_str = env_opt_str("TRITON_MOCK_PTX_VERSION") + dump_ptxas_log: env_bool = env_bool("TRITON_DUMP_PTXAS_LOG") + + libdevice_path: env_opt_str = env_opt_str("TRITON_LIBDEVICE_PATH") + libcuda_path: env_opt_str = env_opt_str("TRITON_LIBCUDA_PATH") + + +class amd_knobs(base_knobs): + use_buffer_ops: env_bool = env_bool("AMDGCN_USE_BUFFER_OPS", True) + # Note: This requires use_buffer_ops be true to have any effect + use_buffer_atomics: env_bool = env_bool("AMDGCN_USE_BUFFER_ATOMICS", True) + # Note: This requires use_buffer_ops be true to have any effect + buffer_ops_analyze_small_tensor_range: env_bool = env_bool("AMDGCN_ANALYZE_SMALL_TENSOR_RANGE", False) + dump_amdgcn: env_bool = env_bool("AMDGCN_ENABLE_DUMP") + libhip_path: env_opt_str = env_opt_str("TRITON_LIBHIP_PATH") + + # We use strs so that we can have a default value based on other runtime info + use_block_pingpong: env_opt_bool = env_opt_bool("TRITON_HIP_USE_BLOCK_PINGPONG") + use_in_thread_transpose: env_opt_bool = env_opt_bool("TRITON_HIP_USE_IN_THREAD_TRANSPOSE") + + use_async_copy: env_bool = env_bool("TRITON_HIP_USE_ASYNC_COPY") + scalarize_packed_fops: env_bool = env_bool("AMDGCN_SCALARIZE_PACKED_FOPS") + + +class proton_knobs(base_knobs): + disable: env_bool = env_bool("TRITON_PROTON_DISABLE", False) + cupti_lib_dir: env_str = env_str( + "TRITON_CUPTI_LIB_PATH", + str(pathlib.Path(__file__).parent.absolute() / "backends" / "nvidia" / "lib" / "cupti")) + enable_nvtx: env_bool = env_bool("TRITON_ENABLE_NVTX", True) + + +build = build_knobs() +redis = redis_knobs() +cache = cache_knobs() +compilation = compilation_knobs() +autotuning = autotuning_knobs() +runtime = runtime_knobs() +language = language_knobs() +nvidia = nvidia_knobs() +amd = amd_knobs() +proton = proton_knobs() + + +def refresh_knobs(): + runtime.debug = env_bool("TRITON_DEBUG").get() + compilation.instrumentation_mode = env_str("TRITON_INSTRUMENTATION_MODE", "").get() diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/triton/testing.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/triton/testing.py new file mode 100644 index 0000000000000000000000000000000000000000..8cfc40a97182239f1b56b850109acaaa61c0e15d --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/triton/testing.py @@ -0,0 +1,542 @@ +import functools +import math +import os +import statistics +import subprocess +import sys +from contextlib import contextmanager +from typing import Any, Dict, List +from . import language as tl +from . import runtime + + +def nvsmi(attrs): + attrs = ','.join(attrs) + cmd = ['nvidia-smi', '-i', '0', '--query-gpu=' + attrs, '--format=csv,noheader,nounits'] + out = subprocess.check_output(cmd) + ret = out.decode(sys.stdout.encoding).split(',') + ret = [int(x) for x in ret] + return ret + + +# pure Python implementation of np.quantile/torch.quantile +# to avoid unnecessary runtime dependency on numpy/torch + + +def _quantile(a, q): + n = len(a) + a = sorted(a) + + def get_quantile(q): + if not (0 <= q <= 1): + raise ValueError("Quantiles must be in the range [0, 1]") + point = q * (n - 1) + lower = math.floor(point) + upper = math.ceil(point) + t = point - lower + return (1 - t) * a[lower] + t * a[upper] + + return [get_quantile(q) for q in q] + + +def _summarize_statistics(times, quantiles, return_mode): + if quantiles is not None: + ret = _quantile(times, quantiles) + if len(ret) == 1: + ret = ret[0] + return ret + if return_mode == "all": + return times + elif return_mode == "min": + return min(times) + elif return_mode == "max": + return max(times) + elif return_mode == "mean": + return statistics.mean(times) + elif return_mode == "median": + return statistics.median(times) + + +def do_bench_cudagraph(fn, rep=20, grad_to_none=None, quantiles=None, return_mode="mean"): + """ + Benchmark the runtime of the provided function. + + :param fn: Function to benchmark + :type fn: Callable + :param rep: Repetition time (in ms) + :type rep: int + :param grad_to_none: Reset the gradient of the provided tensor to None + :type grad_to_none: torch.tensor, optional + :param return_mode: The statistical measure to return. Options are "min", "max", "mean", "median", or "all". Default is "mean". + :type return_mode: str + """ + import torch + assert return_mode in ["min", "max", "mean", "median", "all"] + + with torch.cuda.stream(torch.cuda.Stream()): + # warmup + fn() + if grad_to_none is not None: + for x in grad_to_none: + x.detach_() + x.requires_grad_(True) + x.grad = None + # step 1 - we estimate the amount of time the kernel call takes + # NOTE: this estimate isn't super accurate because the GPU isn't warmed up at this point + # but it is probably good enough + # NOTE: we don't use a graph to estimate the runtime because creating a graph is expensive, + # ~300ms on A100, so we default to the same method used in `do_bench` (minus the L2 + # cache flush). + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + start_event.record() + for _ in range(5): + fn() + end_event.record() + torch.cuda.synchronize() + estimate_ms = start_event.elapsed_time(end_event) / 5 + # Rewrite to avoid possible division by 0 issues with fast benchmarks + if estimate_ms == 0: + n_repeat = 1000 + else: + n_repeat = max(1, int(rep / estimate_ms)) + # step 2 - construct a cuda graph with `n_repeat` unrolled function calls to minimize + # host overhead + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + for _ in range(n_repeat): + if grad_to_none is not None: + for x in grad_to_none: + x.grad = None + fn() + torch.cuda.synchronize() + # measure time and return + ret = [] + n_retries = 10 + for _ in range(n_retries): + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + start_event.record() + g.replay() + end_event.record() + torch.cuda.synchronize() + ret += [start_event.elapsed_time(end_event) / n_repeat] + return _summarize_statistics(ret, quantiles, return_mode) + + +def do_bench(fn, warmup=25, rep=100, grad_to_none=None, quantiles=None, return_mode="mean"): + """ + Benchmark the runtime of the provided function. By default, return the median runtime of :code:`fn` along with + the 20-th and 80-th performance percentile. + + :param fn: Function to benchmark + :type fn: Callable + :param warmup: Warmup time (in ms) + :type warmup: int + :param rep: Repetition time (in ms) + :type rep: int + :param grad_to_none: Reset the gradient of the provided tensor to None + :type grad_to_none: torch.tensor, optional + :param quantiles: Performance percentile to return in addition to the median. + :type quantiles: list[float], optional + :param return_mode: The statistical measure to return. Options are "min", "max", "mean", "median", or "all". Default is "mean". + :type return_mode: str + """ + assert return_mode in ["min", "max", "mean", "median", "all"] + + di = runtime.driver.active.get_device_interface() + + fn() + di.synchronize() + + cache = runtime.driver.active.get_empty_cache_for_benchmark() + + # Estimate the runtime of the function + start_event = di.Event(enable_timing=True) + end_event = di.Event(enable_timing=True) + start_event.record() + for _ in range(5): + runtime.driver.active.clear_cache(cache) + fn() + end_event.record() + di.synchronize() + estimate_ms = start_event.elapsed_time(end_event) / 5 + + # compute number of warmup and repeat + n_warmup = max(1, int(warmup / estimate_ms)) + n_repeat = max(1, int(rep / estimate_ms)) + start_event = [di.Event(enable_timing=True) for i in range(n_repeat)] + end_event = [di.Event(enable_timing=True) for i in range(n_repeat)] + # Warm-up + for _ in range(n_warmup): + fn() + # Benchmark + for i in range(n_repeat): + # we don't want `fn` to accumulate gradient values + # if it contains a backward pass. So we clear the + # provided gradients + if grad_to_none is not None: + for x in grad_to_none: + x.grad = None + # we clear the L2 cache before each run + runtime.driver.active.clear_cache(cache) + # record time of `fn` + start_event[i].record() + fn() + end_event[i].record() + # Record clocks + di.synchronize() + times = [s.elapsed_time(e) for s, e in zip(start_event, end_event)] + return _summarize_statistics(times, quantiles, return_mode) + + +def assert_close(x, y, atol=None, rtol=None, err_msg=''): + """ + Asserts that two inputs are close within a certain tolerance. + + :param x: The first input. + :type x: scala, list, numpy.ndarray, or torch.Tensor + :param y: The second input. + :type y: scala, list, numpy.ndarray, or torch.Tensor + :param atol: The absolute tolerance. Default value is 1e-2. + :type atol: float, optional + :param rtol: The relative tolerance. Default value is 0. + :type rtol: float, optional + :param err_msg: The error message to use if the assertion fails. + :type err_msg: str + """ + import numpy as np + import torch + + # canonicalize arguments to be tensors + if not isinstance(x, torch.Tensor): + x = torch.tensor(x) + if not isinstance(y, torch.Tensor): + y = torch.tensor(y) + # absolute tolerance + if atol is None: + atol = 1e-2 + atol = atol(x.dtype) if callable(atol) else atol + # relative tolerance hook + if rtol is None: + rtol = 0. + rtol = rtol(x.dtype) if callable(rtol) else rtol + # we use numpy instead of pytorch + # as it seems more memory efficient + # pytorch tends to oom on large tensors + if isinstance(x, torch.Tensor): + if x.dtype == torch.bfloat16: + x = x.float() + x = x.cpu().detach().numpy() + if isinstance(y, torch.Tensor): + if y.dtype == torch.bfloat16: + y = y.float() + y = y.cpu().detach().numpy() + # we handle size==1 case separately as we can + # provide better error message there + if x.size > 1 or y.size > 1: + np.testing.assert_allclose(x, y, atol=atol, rtol=rtol, equal_nan=True) + return + if not np.allclose(x, y, atol=atol, rtol=rtol): + raise AssertionError(f'{err_msg} {x} is not close to {y} (atol={atol}, rtol={rtol})') + + +class Benchmark: + """ + This class is used by the :code:`perf_report` function to generate line plots with a concise API. + """ + + def __init__( + self, + x_names: List[str], + x_vals: List[Any], + line_arg: str, + line_vals: List[Any], + line_names: List[str], + plot_name: str, + args: Dict[str, Any], + xlabel: str = '', + ylabel: str = '', + x_log: bool = False, + y_log: bool = False, + styles=None, + ): + """ + Constructor. + x_vals can be a list of scalars or a list of tuples/lists. If x_vals is a list + of scalars and there are multiple x_names, all arguments will have the same value. + If x_vals is a list of tuples/lists, each element should have the same length as + x_names. + + :param x_names: Name of the arguments that should appear on the x axis of the plot. + :type x_names: List[str] + :param x_vals: List of values to use for the arguments in :code:`x_names`. + :type x_vals: List[Any] + :param line_arg: Argument name for which different values correspond to different lines in the plot. + :type line_arg: str + :param line_vals: List of values to use for the arguments in :code:`line_arg`. + :type line_vals: List[Any] + :param line_names: Label names for the different lines. + :type line_names: List[str] + :param plot_name: Name of the plot. + :type plot_name: str + :param args: Dictionary of keyword arguments to remain fixed throughout the benchmark. + :type args: Dict[str, Any] + :param xlabel: Label for the x axis of the plot. + :type xlabel: str, optional + :param ylabel: Label for the y axis of the plot. + :type ylabel: str, optional + :param x_log: Whether the x axis should be log scale. + :type x_log: bool, optional + :param y_log: Whether the y axis should be log scale. + :type y_log: bool, optional + :param styles: A list of tuples, where each tuple contains two elements: a color and a linestyle. + :type styles: list[tuple[str, str]] + """ + self.x_names = x_names + self.x_vals = x_vals + self.x_log = x_log + self.line_arg = line_arg + self.line_vals = line_vals + self.line_names = line_names + self.y_log = y_log + self.styles = styles + # plot info + self.xlabel = xlabel + self.ylabel = ylabel + self.plot_name = plot_name + self.args = args + + +class Mark: + + def __init__(self, fn, benchmarks): + self.fn = fn + self.benchmarks = benchmarks + + def _run(self, bench: Benchmark, save_path: str, show_plots: bool, print_data: bool, diff_col=False, + save_precision=6, **kwrags): + import os + + import matplotlib.pyplot as plt + import pandas as pd + y_mean_labels = [f'{x} ({bench.ylabel})' for x in bench.line_names] + y_min_labels = [f'{x}-min ({bench.ylabel})' for x in bench.line_names] + y_max_labels = [f'{x}-max ({bench.ylabel})' for x in bench.line_names] + x_names = list(bench.x_names) + df = pd.DataFrame(columns=x_names + y_mean_labels + y_min_labels + y_max_labels) + for x in bench.x_vals: + # x can be a single value or a sequence of values. + if not isinstance(x, (list, tuple)): + x = [x for _ in x_names] + + if len(x) != len(x_names): + raise ValueError(f"Expected {len(x_names)} values, got {x}") + x_args = dict(zip(x_names, x)) + + row_mean, row_min, row_max = [], [], [] + for y in bench.line_vals: + ret = self.fn(**x_args, **{bench.line_arg: y}, **bench.args, **kwrags) + try: + y_mean, y_min, y_max = ret + except TypeError: + y_mean, y_min, y_max = ret, None, None + row_mean += [y_mean] + row_min += [y_min] + row_max += [y_max] + df.loc[len(df)] = list(x) + row_mean + row_min + row_max + + if bench.plot_name: + plt.figure() + ax = plt.subplot() + # Plot first x value on x axis if there are multiple. + first_x = x_names[0] + for i, (mean_label, min_label, max_label) in enumerate(zip(y_mean_labels, y_min_labels, y_max_labels)): + y_min, y_max = df[min_label], df[max_label] + col = bench.styles[i][0] if bench.styles else None + sty = bench.styles[i][1] if bench.styles else None + ax.plot(df[first_x], df[mean_label], label=mean_label, color=col, ls=sty) + if not y_min.isnull().all() and not y_max.isnull().all(): + y_min = y_min.astype(float) + y_max = y_max.astype(float) + ax.fill_between(df[first_x], y_min, y_max, alpha=0.15, color=col) + ax.legend() + ax.set_xlabel(bench.xlabel or first_x) + ax.set_ylabel(bench.ylabel) + # ax.set_title(bench.plot_name) + ax.set_xscale("log" if bench.x_log else "linear") + ax.set_yscale("log" if bench.y_log else "linear") + if show_plots: + plt.show() + if save_path: + plt.savefig(os.path.join(save_path, f"{bench.plot_name}.png")) + df = df[x_names + y_mean_labels] + if diff_col and df.shape[1] == 2: + col0, col1 = df.columns.tolist() + df['Diff'] = df[col1] - df[col0] + + if print_data: + print(bench.plot_name + ':') + print(df.to_string()) + if save_path: + df.to_csv(os.path.join(save_path, f"{bench.plot_name}.csv"), float_format=f"%.{save_precision}f", + index=False) + return df + + def run(self, show_plots=False, print_data=False, save_path='', return_df=False, **kwargs): + has_single_bench = isinstance(self.benchmarks, Benchmark) + benchmarks = [self.benchmarks] if has_single_bench else self.benchmarks + result_dfs = [] + try: + for bench in benchmarks: + result_dfs.append(self._run(bench, save_path, show_plots, print_data, **kwargs)) + finally: + if save_path: + # Create directory if it doesn't exist + os.makedirs(save_path, exist_ok=True) + with open(os.path.join(save_path, "results.html"), "w") as html: + html.write("\n") + for bench in benchmarks[:len(result_dfs)]: + html.write(f"\n") + html.write("\n") + if return_df: + if has_single_bench: + return result_dfs[0] + else: + return result_dfs + return None + + +def perf_report(benchmarks): + """ + Mark a function for benchmarking. The benchmark can then be executed by using the :code:`.run` method on the return value. + + :param benchmarks: Benchmarking configurations. + :type benchmarks: List of :class:`Benchmark` + """ + wrapper = lambda fn: Mark(fn, benchmarks) + return wrapper + + +def get_dram_gbps(device=None): + ''' return DRAM bandwidth in GB/s ''' + + from .runtime import driver + if device is None: + device = driver.active.get_device_interface().current_device() + mem_clock_khz = driver.active.utils.get_device_properties(device)["mem_clock_rate"] # in kHz + bus_width = driver.active.utils.get_device_properties(device)["mem_bus_width"] + bw_gbps = mem_clock_khz * bus_width * 2 / 1e6 / 8 # In GB/s + return bw_gbps + + +def get_max_tensorcore_tflops(dtype, clock_rate, device=None): + import torch + + from .runtime import driver + if not device: + device = torch.cuda.current_device() + + num_subcores = driver.active.utils.get_device_properties(device)["multiprocessor_count"] * 4 + capability = torch.cuda.get_device_capability(device) + if capability[0] < 8: + assert dtype == torch.float16 + ops_per_sub_core = 256 # 2 4x4x4 Tensor Cores + else: + if dtype in [torch.float32, torch.int32]: + ops_per_sub_core = 256 + elif dtype in [torch.float16, torch.bfloat16, torch.int16]: + ops_per_sub_core = 512 + elif dtype in [torch.int8, tl.float8e4nv, tl.float8e4b15, tl.float8e5]: + ops_per_sub_core = 1024 + else: + raise RuntimeError("dtype not supported") + tflops = num_subcores * clock_rate * ops_per_sub_core * 1e-9 + return tflops + + +# create decorator that wraps test function into +# a cuda-memcheck system call + + +def cuda_memcheck(**target_kwargs): + + def decorator(test_fn): + + @functools.wraps(test_fn) + def wrapper(*args, **kwargs): + import psutil + ppid_name = psutil.Process(os.getppid()).name() + run_cuda_memcheck = target_kwargs.items() <= kwargs.items() + if run_cuda_memcheck and ppid_name != "cuda-memcheck": + path = os.path.realpath(test_fn.__globals__["__file__"]) + # get path of current file + env = {"PATH": os.environ["PATH"], "PYTORCH_NO_CUDA_MEMORY_CACHING": "1"} + assert 'request' in kwargs, "memcheck'ed test must have a (possibly unused) `request` fixture" + test_id = kwargs['request'].node.callspec.id + cmd = f"{path}::{test_fn.__name__}[{test_id}]" + out = subprocess.run(["cuda-memcheck", "pytest", "-vs", cmd], capture_output=True, env=env) + assert out.returncode == 0, "cuda-memcheck returned an error: bounds checking failed" + assert "ERROR SUMMARY: 0 errors" in str(out.stdout) + else: + test_fn(*args, **kwargs) + + return wrapper + + return decorator + + +@contextmanager +def set_gpu_clock(ref_sm_clock=1350, ref_mem_clock=1215): + try: + subprocess.check_output(["nvidia-smi", "-i", "0", "-pm", "1"]) + subprocess.check_output([ + "nvidia-smi", + "-i", + "0", + f"--lock-gpu-clocks={ref_sm_clock},{ref_sm_clock}", + ]) + subprocess.check_output([ + "nvidia-smi", + "-i", + "0", + f"--lock-memory-clocks={ref_mem_clock},{ref_mem_clock}", + ]) + cur_sm_clock = nvsmi(["clocks.current.sm"])[0] + cur_mem_clock = nvsmi(["clocks.current.memory"])[0] + assert abs(cur_sm_clock - ref_sm_clock) < 10, f"GPU SMs must run at {ref_sm_clock} MHz" + assert abs(cur_mem_clock - ref_mem_clock) < 10, f"GPU SMs must run at {ref_mem_clock} MHz" + tflops = 1e-6 * 2 * 108 * 4 * 256 * ref_sm_clock + gbps = 640 * 2 * ref_mem_clock * 1e-3 + yield tflops, gbps + finally: + subprocess.check_output(["nvidia-smi", "-i", "0", "-pm", "0"]) + subprocess.check_output(["nvidia-smi", "-i", "0", "-rgc"]) + subprocess.check_output(["nvidia-smi", "-i", "0", "-rmc"]) + + +def get_max_simd_tflops(dtype, clock_rate, device=None): + import torch + + from .runtime import driver + if not device: + device = torch.cuda.current_device() + + num_subcores = driver.active.utils.get_device_properties(device)["multiprocessor_count"] * 4 + capability = torch.cuda.get_device_capability() + if capability[0] < 8: + if dtype == torch.float32: + ops_per_sub_core = 32 # 2*16 + elif dtype == torch.float16: + ops_per_sub_core = 64 + else: + raise RuntimeError("dtype not supported") + else: + if dtype == torch.float32: + ops_per_sub_core = 32 + elif dtype in [torch.float16, torch.bfloat16]: + ops_per_sub_core = 64 + else: + raise RuntimeError("dtype not supported") + tflops = num_subcores * clock_rate * ops_per_sub_core * 1e-9 + return tflops diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/wheel/__init__.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/wheel/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3ab8f72d8b2e31b1cb8456ee1c4b8e62d4aba252 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/wheel/__init__.py @@ -0,0 +1,3 @@ +from __future__ import annotations + +__version__ = "0.45.1" diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/wheel/__main__.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/wheel/__main__.py new file mode 100644 index 0000000000000000000000000000000000000000..0be74537494dc2cf18c2e3b318ffd22b886aef6b --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/wheel/__main__.py @@ -0,0 +1,23 @@ +""" +Wheel command line tool (enable python -m wheel syntax) +""" + +from __future__ import annotations + +import sys + + +def main(): # needed for console script + if __package__ == "": + # To be able to run 'python wheel-0.9.whl/wheel': + import os.path + + path = os.path.dirname(os.path.dirname(__file__)) + sys.path[0:0] = [path] + import wheel.cli + + sys.exit(wheel.cli.main()) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/wheel/_bdist_wheel.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/wheel/_bdist_wheel.py new file mode 100644 index 0000000000000000000000000000000000000000..88973ebfb88f9521e5f0613b5cedf1204d10a8f1 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/wheel/_bdist_wheel.py @@ -0,0 +1,613 @@ +""" +Create a wheel (.whl) distribution. + +A wheel is a built archive format. +""" + +from __future__ import annotations + +import os +import re +import shutil +import stat +import struct +import sys +import sysconfig +import warnings +from email.generator import BytesGenerator, Generator +from email.policy import EmailPolicy +from glob import iglob +from shutil import rmtree +from typing import TYPE_CHECKING, Callable, Iterable, Literal, Sequence, cast +from zipfile import ZIP_DEFLATED, ZIP_STORED + +import setuptools +from setuptools import Command + +from . import __version__ as wheel_version +from .metadata import pkginfo_to_metadata +from .util import log +from .vendored.packaging import tags +from .vendored.packaging import version as _packaging_version +from .wheelfile import WheelFile + +if TYPE_CHECKING: + import types + +# ensure Python logging is configured +try: + __import__("setuptools.logging") +except ImportError: + # setuptools < ?? + from . import _setuptools_logging + + _setuptools_logging.configure() + + +def safe_name(name: str) -> str: + """Convert an arbitrary string to a standard distribution name + Any runs of non-alphanumeric/. characters are replaced with a single '-'. + """ + return re.sub("[^A-Za-z0-9.]+", "-", name) + + +def safe_version(version: str) -> str: + """ + Convert an arbitrary string to a standard version string + """ + try: + # normalize the version + return str(_packaging_version.Version(version)) + except _packaging_version.InvalidVersion: + version = version.replace(" ", ".") + return re.sub("[^A-Za-z0-9.]+", "-", version) + + +setuptools_major_version = int(setuptools.__version__.split(".")[0]) + +PY_LIMITED_API_PATTERN = r"cp3\d" + + +def _is_32bit_interpreter() -> bool: + return struct.calcsize("P") == 4 + + +def python_tag() -> str: + return f"py{sys.version_info[0]}" + + +def get_platform(archive_root: str | None) -> str: + """Return our platform name 'win32', 'linux_x86_64'""" + result = sysconfig.get_platform() + if result.startswith("macosx") and archive_root is not None: + from .macosx_libfile import calculate_macosx_platform_tag + + result = calculate_macosx_platform_tag(archive_root, result) + elif _is_32bit_interpreter(): + if result == "linux-x86_64": + # pip pull request #3497 + result = "linux-i686" + elif result == "linux-aarch64": + # packaging pull request #234 + # TODO armv8l, packaging pull request #690 => this did not land + # in pip/packaging yet + result = "linux-armv7l" + + return result.replace("-", "_") + + +def get_flag( + var: str, fallback: bool, expected: bool = True, warn: bool = True +) -> bool: + """Use a fallback value for determining SOABI flags if the needed config + var is unset or unavailable.""" + val = sysconfig.get_config_var(var) + if val is None: + if warn: + warnings.warn( + f"Config variable '{var}' is unset, Python ABI tag may be incorrect", + RuntimeWarning, + stacklevel=2, + ) + return fallback + return val == expected + + +def get_abi_tag() -> str | None: + """Return the ABI tag based on SOABI (if available) or emulate SOABI (PyPy2).""" + soabi: str = sysconfig.get_config_var("SOABI") + impl = tags.interpreter_name() + if not soabi and impl in ("cp", "pp") and hasattr(sys, "maxunicode"): + d = "" + m = "" + u = "" + if get_flag("Py_DEBUG", hasattr(sys, "gettotalrefcount"), warn=(impl == "cp")): + d = "d" + + if get_flag( + "WITH_PYMALLOC", + impl == "cp", + warn=(impl == "cp" and sys.version_info < (3, 8)), + ) and sys.version_info < (3, 8): + m = "m" + + abi = f"{impl}{tags.interpreter_version()}{d}{m}{u}" + elif soabi and impl == "cp" and soabi.startswith("cpython"): + # non-Windows + abi = "cp" + soabi.split("-")[1] + elif soabi and impl == "cp" and soabi.startswith("cp"): + # Windows + abi = soabi.split("-")[0] + elif soabi and impl == "pp": + # we want something like pypy36-pp73 + abi = "-".join(soabi.split("-")[:2]) + abi = abi.replace(".", "_").replace("-", "_") + elif soabi and impl == "graalpy": + abi = "-".join(soabi.split("-")[:3]) + abi = abi.replace(".", "_").replace("-", "_") + elif soabi: + abi = soabi.replace(".", "_").replace("-", "_") + else: + abi = None + + return abi + + +def safer_name(name: str) -> str: + return safe_name(name).replace("-", "_") + + +def safer_version(version: str) -> str: + return safe_version(version).replace("-", "_") + + +def remove_readonly( + func: Callable[..., object], + path: str, + excinfo: tuple[type[Exception], Exception, types.TracebackType], +) -> None: + remove_readonly_exc(func, path, excinfo[1]) + + +def remove_readonly_exc(func: Callable[..., object], path: str, exc: Exception) -> None: + os.chmod(path, stat.S_IWRITE) + func(path) + + +class bdist_wheel(Command): + description = "create a wheel distribution" + + supported_compressions = { + "stored": ZIP_STORED, + "deflated": ZIP_DEFLATED, + } + + user_options = [ + ("bdist-dir=", "b", "temporary directory for creating the distribution"), + ( + "plat-name=", + "p", + "platform name to embed in generated filenames " + f"(default: {get_platform(None)})", + ), + ( + "keep-temp", + "k", + "keep the pseudo-installation tree around after " + "creating the distribution archive", + ), + ("dist-dir=", "d", "directory to put final built distributions in"), + ("skip-build", None, "skip rebuilding everything (for testing/debugging)"), + ( + "relative", + None, + "build the archive using relative paths (default: false)", + ), + ( + "owner=", + "u", + "Owner name used when creating a tar file [default: current user]", + ), + ( + "group=", + "g", + "Group name used when creating a tar file [default: current group]", + ), + ("universal", None, "make a universal wheel (default: false)"), + ( + "compression=", + None, + "zipfile compression (one of: {}) (default: 'deflated')".format( + ", ".join(supported_compressions) + ), + ), + ( + "python-tag=", + None, + f"Python implementation compatibility tag (default: '{python_tag()}')", + ), + ( + "build-number=", + None, + "Build number for this particular version. " + "As specified in PEP-0427, this must start with a digit. " + "[default: None]", + ), + ( + "py-limited-api=", + None, + "Python tag (cp32|cp33|cpNN) for abi3 wheel tag (default: false)", + ), + ] + + boolean_options = ["keep-temp", "skip-build", "relative", "universal"] + + def initialize_options(self): + self.bdist_dir: str = None + self.data_dir = None + self.plat_name: str | None = None + self.plat_tag = None + self.format = "zip" + self.keep_temp = False + self.dist_dir: str | None = None + self.egginfo_dir = None + self.root_is_pure: bool | None = None + self.skip_build = None + self.relative = False + self.owner = None + self.group = None + self.universal: bool = False + self.compression: str | int = "deflated" + self.python_tag: str = python_tag() + self.build_number: str | None = None + self.py_limited_api: str | Literal[False] = False + self.plat_name_supplied = False + + def finalize_options(self): + if self.bdist_dir is None: + bdist_base = self.get_finalized_command("bdist").bdist_base + self.bdist_dir = os.path.join(bdist_base, "wheel") + + egg_info = self.distribution.get_command_obj("egg_info") + egg_info.ensure_finalized() # needed for correct `wheel_dist_name` + + self.data_dir = self.wheel_dist_name + ".data" + self.plat_name_supplied = self.plat_name is not None + + try: + self.compression = self.supported_compressions[self.compression] + except KeyError: + raise ValueError(f"Unsupported compression: {self.compression}") from None + + need_options = ("dist_dir", "plat_name", "skip_build") + + self.set_undefined_options("bdist", *zip(need_options, need_options)) + + self.root_is_pure = not ( + self.distribution.has_ext_modules() or self.distribution.has_c_libraries() + ) + + if self.py_limited_api and not re.match( + PY_LIMITED_API_PATTERN, self.py_limited_api + ): + raise ValueError(f"py-limited-api must match '{PY_LIMITED_API_PATTERN}'") + + # Support legacy [wheel] section for setting universal + wheel = self.distribution.get_option_dict("wheel") + if "universal" in wheel: + # please don't define this in your global configs + log.warning( + "The [wheel] section is deprecated. Use [bdist_wheel] instead.", + ) + val = wheel["universal"][1].strip() + if val.lower() in ("1", "true", "yes"): + self.universal = True + + if self.build_number is not None and not self.build_number[:1].isdigit(): + raise ValueError("Build tag (build-number) must start with a digit.") + + @property + def wheel_dist_name(self): + """Return distribution full name with - replaced with _""" + components = ( + safer_name(self.distribution.get_name()), + safer_version(self.distribution.get_version()), + ) + if self.build_number: + components += (self.build_number,) + return "-".join(components) + + def get_tag(self) -> tuple[str, str, str]: + # bdist sets self.plat_name if unset, we should only use it for purepy + # wheels if the user supplied it. + if self.plat_name_supplied: + plat_name = cast(str, self.plat_name) + elif self.root_is_pure: + plat_name = "any" + else: + # macosx contains system version in platform name so need special handle + if self.plat_name and not self.plat_name.startswith("macosx"): + plat_name = self.plat_name + else: + # on macosx always limit the platform name to comply with any + # c-extension modules in bdist_dir, since the user can specify + # a higher MACOSX_DEPLOYMENT_TARGET via tools like CMake + + # on other platforms, and on macosx if there are no c-extension + # modules, use the default platform name. + plat_name = get_platform(self.bdist_dir) + + if _is_32bit_interpreter(): + if plat_name in ("linux-x86_64", "linux_x86_64"): + plat_name = "linux_i686" + if plat_name in ("linux-aarch64", "linux_aarch64"): + # TODO armv8l, packaging pull request #690 => this did not land + # in pip/packaging yet + plat_name = "linux_armv7l" + + plat_name = ( + plat_name.lower().replace("-", "_").replace(".", "_").replace(" ", "_") + ) + + if self.root_is_pure: + if self.universal: + impl = "py2.py3" + else: + impl = self.python_tag + tag = (impl, "none", plat_name) + else: + impl_name = tags.interpreter_name() + impl_ver = tags.interpreter_version() + impl = impl_name + impl_ver + # We don't work on CPython 3.1, 3.0. + if self.py_limited_api and (impl_name + impl_ver).startswith("cp3"): + impl = self.py_limited_api + abi_tag = "abi3" + else: + abi_tag = str(get_abi_tag()).lower() + tag = (impl, abi_tag, plat_name) + # issue gh-374: allow overriding plat_name + supported_tags = [ + (t.interpreter, t.abi, plat_name) for t in tags.sys_tags() + ] + assert ( + tag in supported_tags + ), f"would build wheel with unsupported tag {tag}" + return tag + + def run(self): + build_scripts = self.reinitialize_command("build_scripts") + build_scripts.executable = "python" + build_scripts.force = True + + build_ext = self.reinitialize_command("build_ext") + build_ext.inplace = False + + if not self.skip_build: + self.run_command("build") + + install = self.reinitialize_command("install", reinit_subcommands=True) + install.root = self.bdist_dir + install.compile = False + install.skip_build = self.skip_build + install.warn_dir = False + + # A wheel without setuptools scripts is more cross-platform. + # Use the (undocumented) `no_ep` option to setuptools' + # install_scripts command to avoid creating entry point scripts. + install_scripts = self.reinitialize_command("install_scripts") + install_scripts.no_ep = True + + # Use a custom scheme for the archive, because we have to decide + # at installation time which scheme to use. + for key in ("headers", "scripts", "data", "purelib", "platlib"): + setattr(install, "install_" + key, os.path.join(self.data_dir, key)) + + basedir_observed = "" + + if os.name == "nt": + # win32 barfs if any of these are ''; could be '.'? + # (distutils.command.install:change_roots bug) + basedir_observed = os.path.normpath(os.path.join(self.data_dir, "..")) + self.install_libbase = self.install_lib = basedir_observed + + setattr( + install, + "install_purelib" if self.root_is_pure else "install_platlib", + basedir_observed, + ) + + log.info(f"installing to {self.bdist_dir}") + + self.run_command("install") + + impl_tag, abi_tag, plat_tag = self.get_tag() + archive_basename = f"{self.wheel_dist_name}-{impl_tag}-{abi_tag}-{plat_tag}" + if not self.relative: + archive_root = self.bdist_dir + else: + archive_root = os.path.join( + self.bdist_dir, self._ensure_relative(install.install_base) + ) + + self.set_undefined_options("install_egg_info", ("target", "egginfo_dir")) + distinfo_dirname = ( + f"{safer_name(self.distribution.get_name())}-" + f"{safer_version(self.distribution.get_version())}.dist-info" + ) + distinfo_dir = os.path.join(self.bdist_dir, distinfo_dirname) + self.egg2dist(self.egginfo_dir, distinfo_dir) + + self.write_wheelfile(distinfo_dir) + + # Make the archive + if not os.path.exists(self.dist_dir): + os.makedirs(self.dist_dir) + + wheel_path = os.path.join(self.dist_dir, archive_basename + ".whl") + with WheelFile(wheel_path, "w", self.compression) as wf: + wf.write_files(archive_root) + + # Add to 'Distribution.dist_files' so that the "upload" command works + getattr(self.distribution, "dist_files", []).append( + ( + "bdist_wheel", + "{}.{}".format(*sys.version_info[:2]), # like 3.7 + wheel_path, + ) + ) + + if not self.keep_temp: + log.info(f"removing {self.bdist_dir}") + if not self.dry_run: + if sys.version_info < (3, 12): + rmtree(self.bdist_dir, onerror=remove_readonly) + else: + rmtree(self.bdist_dir, onexc=remove_readonly_exc) + + def write_wheelfile( + self, wheelfile_base: str, generator: str = f"bdist_wheel ({wheel_version})" + ): + from email.message import Message + + msg = Message() + msg["Wheel-Version"] = "1.0" # of the spec + msg["Generator"] = generator + msg["Root-Is-Purelib"] = str(self.root_is_pure).lower() + if self.build_number is not None: + msg["Build"] = self.build_number + + # Doesn't work for bdist_wininst + impl_tag, abi_tag, plat_tag = self.get_tag() + for impl in impl_tag.split("."): + for abi in abi_tag.split("."): + for plat in plat_tag.split("."): + msg["Tag"] = "-".join((impl, abi, plat)) + + wheelfile_path = os.path.join(wheelfile_base, "WHEEL") + log.info(f"creating {wheelfile_path}") + with open(wheelfile_path, "wb") as f: + BytesGenerator(f, maxheaderlen=0).flatten(msg) + + def _ensure_relative(self, path: str) -> str: + # copied from dir_util, deleted + drive, path = os.path.splitdrive(path) + if path[0:1] == os.sep: + path = drive + path[1:] + return path + + @property + def license_paths(self) -> Iterable[str]: + if setuptools_major_version >= 57: + # Setuptools has resolved any patterns to actual file names + return self.distribution.metadata.license_files or () + + files: set[str] = set() + metadata = self.distribution.get_option_dict("metadata") + if setuptools_major_version >= 42: + # Setuptools recognizes the license_files option but does not do globbing + patterns = cast(Sequence[str], self.distribution.metadata.license_files) + else: + # Prior to those, wheel is entirely responsible for handling license files + if "license_files" in metadata: + patterns = metadata["license_files"][1].split() + else: + patterns = () + + if "license_file" in metadata: + warnings.warn( + 'The "license_file" option is deprecated. Use "license_files" instead.', + DeprecationWarning, + stacklevel=2, + ) + files.add(metadata["license_file"][1]) + + if not files and not patterns and not isinstance(patterns, list): + patterns = ("LICEN[CS]E*", "COPYING*", "NOTICE*", "AUTHORS*") + + for pattern in patterns: + for path in iglob(pattern): + if path.endswith("~"): + log.debug( + f'ignoring license file "{path}" as it looks like a backup' + ) + continue + + if path not in files and os.path.isfile(path): + log.info( + f'adding license file "{path}" (matched pattern "{pattern}")' + ) + files.add(path) + + return files + + def egg2dist(self, egginfo_path: str, distinfo_path: str): + """Convert an .egg-info directory into a .dist-info directory""" + + def adios(p: str) -> None: + """Appropriately delete directory, file or link.""" + if os.path.exists(p) and not os.path.islink(p) and os.path.isdir(p): + shutil.rmtree(p) + elif os.path.exists(p): + os.unlink(p) + + adios(distinfo_path) + + if not os.path.exists(egginfo_path): + # There is no egg-info. This is probably because the egg-info + # file/directory is not named matching the distribution name used + # to name the archive file. Check for this case and report + # accordingly. + import glob + + pat = os.path.join(os.path.dirname(egginfo_path), "*.egg-info") + possible = glob.glob(pat) + err = f"Egg metadata expected at {egginfo_path} but not found" + if possible: + alt = os.path.basename(possible[0]) + err += f" ({alt} found - possible misnamed archive file?)" + + raise ValueError(err) + + if os.path.isfile(egginfo_path): + # .egg-info is a single file + pkg_info = pkginfo_to_metadata(egginfo_path, egginfo_path) + os.mkdir(distinfo_path) + else: + # .egg-info is a directory + pkginfo_path = os.path.join(egginfo_path, "PKG-INFO") + pkg_info = pkginfo_to_metadata(egginfo_path, pkginfo_path) + + # ignore common egg metadata that is useless to wheel + shutil.copytree( + egginfo_path, + distinfo_path, + ignore=lambda x, y: { + "PKG-INFO", + "requires.txt", + "SOURCES.txt", + "not-zip-safe", + }, + ) + + # delete dependency_links if it is only whitespace + dependency_links_path = os.path.join(distinfo_path, "dependency_links.txt") + with open(dependency_links_path, encoding="utf-8") as dependency_links_file: + dependency_links = dependency_links_file.read().strip() + if not dependency_links: + adios(dependency_links_path) + + pkg_info_path = os.path.join(distinfo_path, "METADATA") + serialization_policy = EmailPolicy( + utf8=True, + mangle_from_=False, + max_line_length=0, + ) + with open(pkg_info_path, "w", encoding="utf-8") as out: + Generator(out, policy=serialization_policy).flatten(pkg_info) + + for license_path in self.license_paths: + filename = os.path.basename(license_path) + shutil.copy(license_path, os.path.join(distinfo_path, filename)) + + adios(egginfo_path) diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/wheel/_setuptools_logging.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/wheel/_setuptools_logging.py new file mode 100644 index 0000000000000000000000000000000000000000..a1a2482ba29ac5290c8f7d7688452ec3faf59332 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/wheel/_setuptools_logging.py @@ -0,0 +1,26 @@ +# copied from setuptools.logging, omitting monkeypatching +from __future__ import annotations + +import logging +import sys + + +def _not_warning(record: logging.LogRecord) -> bool: + return record.levelno < logging.WARNING + + +def configure() -> None: + """ + Configure logging to emit warning and above to stderr + and everything else to stdout. This behavior is provided + for compatibility with distutils.log but may change in + the future. + """ + err_handler = logging.StreamHandler() + err_handler.setLevel(logging.WARNING) + out_handler = logging.StreamHandler(sys.stdout) + out_handler.addFilter(_not_warning) + handlers = err_handler, out_handler + logging.basicConfig( + format="{message}", style="{", handlers=handlers, level=logging.DEBUG + ) diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/wheel/bdist_wheel.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/wheel/bdist_wheel.py new file mode 100644 index 0000000000000000000000000000000000000000..dd7b8629e5087fe35a2ea7ab6d3cbbb4cacf0031 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/wheel/bdist_wheel.py @@ -0,0 +1,26 @@ +from typing import TYPE_CHECKING +from warnings import warn + +warn( + "The 'wheel' package is no longer the canonical location of the 'bdist_wheel' " + "command, and will be removed in a future release. Please update to setuptools " + "v70.1 or later which contains an integrated version of this command.", + DeprecationWarning, + stacklevel=1, +) + +if TYPE_CHECKING: + from ._bdist_wheel import bdist_wheel as bdist_wheel +else: + try: + # Better integration/compatibility with setuptools: + # in the case new fixes or PEPs are implemented in setuptools + # there is no need to backport them to the deprecated code base. + # This is useful in the case of old packages in the ecosystem + # that are still used but have low maintenance. + from setuptools.command.bdist_wheel import bdist_wheel + except ImportError: + # Only used in the case of old setuptools versions. + # If the user wants to get the latest fixes/PEPs, + # they are encouraged to address the deprecation warning. + from ._bdist_wheel import bdist_wheel as bdist_wheel diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/wheel/macosx_libfile.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/wheel/macosx_libfile.py new file mode 100644 index 0000000000000000000000000000000000000000..abdfc9eda1d4c1f33270c155e1610fe73bd54263 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/wheel/macosx_libfile.py @@ -0,0 +1,482 @@ +""" +This module contains function to analyse dynamic library +headers to extract system information + +Currently only for MacOSX + +Library file on macosx system starts with Mach-O or Fat field. +This can be distinguish by first 32 bites and it is called magic number. +Proper value of magic number is with suffix _MAGIC. Suffix _CIGAM means +reversed bytes order. +Both fields can occur in two types: 32 and 64 bytes. + +FAT field inform that this library contains few version of library +(typically for different types version). It contains +information where Mach-O headers starts. + +Each section started with Mach-O header contains one library +(So if file starts with this field it contains only one version). + +After filed Mach-O there are section fields. +Each of them starts with two fields: +cmd - magic number for this command +cmdsize - total size occupied by this section information. + +In this case only sections LC_VERSION_MIN_MACOSX (for macosx 10.13 and earlier) +and LC_BUILD_VERSION (for macosx 10.14 and newer) are interesting, +because them contains information about minimal system version. + +Important remarks: +- For fat files this implementation looks for maximum number version. + It not check if it is 32 or 64 and do not compare it with currently built package. + So it is possible to false report higher version that needed. +- All structures signatures are taken form macosx header files. +- I think that binary format will be more stable than `otool` output. + and if apple introduce some changes both implementation will need to be updated. +- The system compile will set the deployment target no lower than + 11.0 for arm64 builds. For "Universal 2" builds use the x86_64 deployment + target when the arm64 target is 11.0. +""" + +from __future__ import annotations + +import ctypes +import os +import sys +from io import BufferedIOBase +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Union + + StrPath = Union[str, os.PathLike[str]] + +"""here the needed const and struct from mach-o header files""" + +FAT_MAGIC = 0xCAFEBABE +FAT_CIGAM = 0xBEBAFECA +FAT_MAGIC_64 = 0xCAFEBABF +FAT_CIGAM_64 = 0xBFBAFECA +MH_MAGIC = 0xFEEDFACE +MH_CIGAM = 0xCEFAEDFE +MH_MAGIC_64 = 0xFEEDFACF +MH_CIGAM_64 = 0xCFFAEDFE + +LC_VERSION_MIN_MACOSX = 0x24 +LC_BUILD_VERSION = 0x32 + +CPU_TYPE_ARM64 = 0x0100000C + +mach_header_fields = [ + ("magic", ctypes.c_uint32), + ("cputype", ctypes.c_int), + ("cpusubtype", ctypes.c_int), + ("filetype", ctypes.c_uint32), + ("ncmds", ctypes.c_uint32), + ("sizeofcmds", ctypes.c_uint32), + ("flags", ctypes.c_uint32), +] +""" +struct mach_header { + uint32_t magic; /* mach magic number identifier */ + cpu_type_t cputype; /* cpu specifier */ + cpu_subtype_t cpusubtype; /* machine specifier */ + uint32_t filetype; /* type of file */ + uint32_t ncmds; /* number of load commands */ + uint32_t sizeofcmds; /* the size of all the load commands */ + uint32_t flags; /* flags */ +}; +typedef integer_t cpu_type_t; +typedef integer_t cpu_subtype_t; +""" + +mach_header_fields_64 = mach_header_fields + [("reserved", ctypes.c_uint32)] +""" +struct mach_header_64 { + uint32_t magic; /* mach magic number identifier */ + cpu_type_t cputype; /* cpu specifier */ + cpu_subtype_t cpusubtype; /* machine specifier */ + uint32_t filetype; /* type of file */ + uint32_t ncmds; /* number of load commands */ + uint32_t sizeofcmds; /* the size of all the load commands */ + uint32_t flags; /* flags */ + uint32_t reserved; /* reserved */ +}; +""" + +fat_header_fields = [("magic", ctypes.c_uint32), ("nfat_arch", ctypes.c_uint32)] +""" +struct fat_header { + uint32_t magic; /* FAT_MAGIC or FAT_MAGIC_64 */ + uint32_t nfat_arch; /* number of structs that follow */ +}; +""" + +fat_arch_fields = [ + ("cputype", ctypes.c_int), + ("cpusubtype", ctypes.c_int), + ("offset", ctypes.c_uint32), + ("size", ctypes.c_uint32), + ("align", ctypes.c_uint32), +] +""" +struct fat_arch { + cpu_type_t cputype; /* cpu specifier (int) */ + cpu_subtype_t cpusubtype; /* machine specifier (int) */ + uint32_t offset; /* file offset to this object file */ + uint32_t size; /* size of this object file */ + uint32_t align; /* alignment as a power of 2 */ +}; +""" + +fat_arch_64_fields = [ + ("cputype", ctypes.c_int), + ("cpusubtype", ctypes.c_int), + ("offset", ctypes.c_uint64), + ("size", ctypes.c_uint64), + ("align", ctypes.c_uint32), + ("reserved", ctypes.c_uint32), +] +""" +struct fat_arch_64 { + cpu_type_t cputype; /* cpu specifier (int) */ + cpu_subtype_t cpusubtype; /* machine specifier (int) */ + uint64_t offset; /* file offset to this object file */ + uint64_t size; /* size of this object file */ + uint32_t align; /* alignment as a power of 2 */ + uint32_t reserved; /* reserved */ +}; +""" + +segment_base_fields = [("cmd", ctypes.c_uint32), ("cmdsize", ctypes.c_uint32)] +"""base for reading segment info""" + +segment_command_fields = [ + ("cmd", ctypes.c_uint32), + ("cmdsize", ctypes.c_uint32), + ("segname", ctypes.c_char * 16), + ("vmaddr", ctypes.c_uint32), + ("vmsize", ctypes.c_uint32), + ("fileoff", ctypes.c_uint32), + ("filesize", ctypes.c_uint32), + ("maxprot", ctypes.c_int), + ("initprot", ctypes.c_int), + ("nsects", ctypes.c_uint32), + ("flags", ctypes.c_uint32), +] +""" +struct segment_command { /* for 32-bit architectures */ + uint32_t cmd; /* LC_SEGMENT */ + uint32_t cmdsize; /* includes sizeof section structs */ + char segname[16]; /* segment name */ + uint32_t vmaddr; /* memory address of this segment */ + uint32_t vmsize; /* memory size of this segment */ + uint32_t fileoff; /* file offset of this segment */ + uint32_t filesize; /* amount to map from the file */ + vm_prot_t maxprot; /* maximum VM protection */ + vm_prot_t initprot; /* initial VM protection */ + uint32_t nsects; /* number of sections in segment */ + uint32_t flags; /* flags */ +}; +typedef int vm_prot_t; +""" + +segment_command_fields_64 = [ + ("cmd", ctypes.c_uint32), + ("cmdsize", ctypes.c_uint32), + ("segname", ctypes.c_char * 16), + ("vmaddr", ctypes.c_uint64), + ("vmsize", ctypes.c_uint64), + ("fileoff", ctypes.c_uint64), + ("filesize", ctypes.c_uint64), + ("maxprot", ctypes.c_int), + ("initprot", ctypes.c_int), + ("nsects", ctypes.c_uint32), + ("flags", ctypes.c_uint32), +] +""" +struct segment_command_64 { /* for 64-bit architectures */ + uint32_t cmd; /* LC_SEGMENT_64 */ + uint32_t cmdsize; /* includes sizeof section_64 structs */ + char segname[16]; /* segment name */ + uint64_t vmaddr; /* memory address of this segment */ + uint64_t vmsize; /* memory size of this segment */ + uint64_t fileoff; /* file offset of this segment */ + uint64_t filesize; /* amount to map from the file */ + vm_prot_t maxprot; /* maximum VM protection */ + vm_prot_t initprot; /* initial VM protection */ + uint32_t nsects; /* number of sections in segment */ + uint32_t flags; /* flags */ +}; +""" + +version_min_command_fields = segment_base_fields + [ + ("version", ctypes.c_uint32), + ("sdk", ctypes.c_uint32), +] +""" +struct version_min_command { + uint32_t cmd; /* LC_VERSION_MIN_MACOSX or + LC_VERSION_MIN_IPHONEOS or + LC_VERSION_MIN_WATCHOS or + LC_VERSION_MIN_TVOS */ + uint32_t cmdsize; /* sizeof(struct min_version_command) */ + uint32_t version; /* X.Y.Z is encoded in nibbles xxxx.yy.zz */ + uint32_t sdk; /* X.Y.Z is encoded in nibbles xxxx.yy.zz */ +}; +""" + +build_version_command_fields = segment_base_fields + [ + ("platform", ctypes.c_uint32), + ("minos", ctypes.c_uint32), + ("sdk", ctypes.c_uint32), + ("ntools", ctypes.c_uint32), +] +""" +struct build_version_command { + uint32_t cmd; /* LC_BUILD_VERSION */ + uint32_t cmdsize; /* sizeof(struct build_version_command) plus */ + /* ntools * sizeof(struct build_tool_version) */ + uint32_t platform; /* platform */ + uint32_t minos; /* X.Y.Z is encoded in nibbles xxxx.yy.zz */ + uint32_t sdk; /* X.Y.Z is encoded in nibbles xxxx.yy.zz */ + uint32_t ntools; /* number of tool entries following this */ +}; +""" + + +def swap32(x: int) -> int: + return ( + ((x << 24) & 0xFF000000) + | ((x << 8) & 0x00FF0000) + | ((x >> 8) & 0x0000FF00) + | ((x >> 24) & 0x000000FF) + ) + + +def get_base_class_and_magic_number( + lib_file: BufferedIOBase, + seek: int | None = None, +) -> tuple[type[ctypes.Structure], int]: + if seek is None: + seek = lib_file.tell() + else: + lib_file.seek(seek) + magic_number = ctypes.c_uint32.from_buffer_copy( + lib_file.read(ctypes.sizeof(ctypes.c_uint32)) + ).value + + # Handle wrong byte order + if magic_number in [FAT_CIGAM, FAT_CIGAM_64, MH_CIGAM, MH_CIGAM_64]: + if sys.byteorder == "little": + BaseClass = ctypes.BigEndianStructure + else: + BaseClass = ctypes.LittleEndianStructure + + magic_number = swap32(magic_number) + else: + BaseClass = ctypes.Structure + + lib_file.seek(seek) + return BaseClass, magic_number + + +def read_data(struct_class: type[ctypes.Structure], lib_file: BufferedIOBase): + return struct_class.from_buffer_copy(lib_file.read(ctypes.sizeof(struct_class))) + + +def extract_macosx_min_system_version(path_to_lib: str): + with open(path_to_lib, "rb") as lib_file: + BaseClass, magic_number = get_base_class_and_magic_number(lib_file, 0) + if magic_number not in [FAT_MAGIC, FAT_MAGIC_64, MH_MAGIC, MH_MAGIC_64]: + return + + if magic_number in [FAT_MAGIC, FAT_CIGAM_64]: + + class FatHeader(BaseClass): + _fields_ = fat_header_fields + + fat_header = read_data(FatHeader, lib_file) + if magic_number == FAT_MAGIC: + + class FatArch(BaseClass): + _fields_ = fat_arch_fields + + else: + + class FatArch(BaseClass): + _fields_ = fat_arch_64_fields + + fat_arch_list = [ + read_data(FatArch, lib_file) for _ in range(fat_header.nfat_arch) + ] + + versions_list: list[tuple[int, int, int]] = [] + for el in fat_arch_list: + try: + version = read_mach_header(lib_file, el.offset) + if version is not None: + if el.cputype == CPU_TYPE_ARM64 and len(fat_arch_list) != 1: + # Xcode will not set the deployment target below 11.0.0 + # for the arm64 architecture. Ignore the arm64 deployment + # in fat binaries when the target is 11.0.0, that way + # the other architectures can select a lower deployment + # target. + # This is safe because there is no arm64 variant for + # macOS 10.15 or earlier. + if version == (11, 0, 0): + continue + versions_list.append(version) + except ValueError: + pass + + if len(versions_list) > 0: + return max(versions_list) + else: + return None + + else: + try: + return read_mach_header(lib_file, 0) + except ValueError: + """when some error during read library files""" + return None + + +def read_mach_header( + lib_file: BufferedIOBase, + seek: int | None = None, +) -> tuple[int, int, int] | None: + """ + This function parses a Mach-O header and extracts + information about the minimal macOS version. + + :param lib_file: reference to opened library file with pointer + """ + base_class, magic_number = get_base_class_and_magic_number(lib_file, seek) + arch = "32" if magic_number == MH_MAGIC else "64" + + class SegmentBase(base_class): + _fields_ = segment_base_fields + + if arch == "32": + + class MachHeader(base_class): + _fields_ = mach_header_fields + + else: + + class MachHeader(base_class): + _fields_ = mach_header_fields_64 + + mach_header = read_data(MachHeader, lib_file) + for _i in range(mach_header.ncmds): + pos = lib_file.tell() + segment_base = read_data(SegmentBase, lib_file) + lib_file.seek(pos) + if segment_base.cmd == LC_VERSION_MIN_MACOSX: + + class VersionMinCommand(base_class): + _fields_ = version_min_command_fields + + version_info = read_data(VersionMinCommand, lib_file) + return parse_version(version_info.version) + elif segment_base.cmd == LC_BUILD_VERSION: + + class VersionBuild(base_class): + _fields_ = build_version_command_fields + + version_info = read_data(VersionBuild, lib_file) + return parse_version(version_info.minos) + else: + lib_file.seek(pos + segment_base.cmdsize) + continue + + +def parse_version(version: int) -> tuple[int, int, int]: + x = (version & 0xFFFF0000) >> 16 + y = (version & 0x0000FF00) >> 8 + z = version & 0x000000FF + return x, y, z + + +def calculate_macosx_platform_tag(archive_root: StrPath, platform_tag: str) -> str: + """ + Calculate proper macosx platform tag basing on files which are included to wheel + + Example platform tag `macosx-10.14-x86_64` + """ + prefix, base_version, suffix = platform_tag.split("-") + base_version = tuple(int(x) for x in base_version.split(".")) + base_version = base_version[:2] + if base_version[0] > 10: + base_version = (base_version[0], 0) + assert len(base_version) == 2 + if "MACOSX_DEPLOYMENT_TARGET" in os.environ: + deploy_target = tuple( + int(x) for x in os.environ["MACOSX_DEPLOYMENT_TARGET"].split(".") + ) + deploy_target = deploy_target[:2] + if deploy_target[0] > 10: + deploy_target = (deploy_target[0], 0) + if deploy_target < base_version: + sys.stderr.write( + "[WARNING] MACOSX_DEPLOYMENT_TARGET is set to a lower value ({}) than " + "the version on which the Python interpreter was compiled ({}), and " + "will be ignored.\n".format( + ".".join(str(x) for x in deploy_target), + ".".join(str(x) for x in base_version), + ) + ) + else: + base_version = deploy_target + + assert len(base_version) == 2 + start_version = base_version + versions_dict: dict[str, tuple[int, int]] = {} + for dirpath, _dirnames, filenames in os.walk(archive_root): + for filename in filenames: + if filename.endswith(".dylib") or filename.endswith(".so"): + lib_path = os.path.join(dirpath, filename) + min_ver = extract_macosx_min_system_version(lib_path) + if min_ver is not None: + min_ver = min_ver[0:2] + if min_ver[0] > 10: + min_ver = (min_ver[0], 0) + versions_dict[lib_path] = min_ver + + if len(versions_dict) > 0: + base_version = max(base_version, max(versions_dict.values())) + + # macosx platform tag do not support minor bugfix release + fin_base_version = "_".join([str(x) for x in base_version]) + if start_version < base_version: + problematic_files = [k for k, v in versions_dict.items() if v > start_version] + problematic_files = "\n".join(problematic_files) + if len(problematic_files) == 1: + files_form = "this file" + else: + files_form = "these files" + error_message = ( + "[WARNING] This wheel needs a higher macOS version than {} " + "To silence this warning, set MACOSX_DEPLOYMENT_TARGET to at least " + + fin_base_version + + " or recreate " + + files_form + + " with lower " + "MACOSX_DEPLOYMENT_TARGET: \n" + problematic_files + ) + + if "MACOSX_DEPLOYMENT_TARGET" in os.environ: + error_message = error_message.format( + "is set in MACOSX_DEPLOYMENT_TARGET variable." + ) + else: + error_message = error_message.format( + "the version your Python interpreter is compiled against." + ) + + sys.stderr.write(error_message) + + platform_tag = prefix + "_" + fin_base_version + "_" + suffix + return platform_tag diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/wheel/metadata.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/wheel/metadata.py new file mode 100644 index 0000000000000000000000000000000000000000..b8098fa8599a1f142ee33ee551b7a874f5f98840 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/wheel/metadata.py @@ -0,0 +1,183 @@ +""" +Tools for converting old- to new-style metadata. +""" + +from __future__ import annotations + +import functools +import itertools +import os.path +import re +import textwrap +from email.message import Message +from email.parser import Parser +from typing import Generator, Iterable, Iterator, Literal + +from .vendored.packaging.requirements import Requirement + + +def _nonblank(str: str) -> bool | Literal[""]: + return str and not str.startswith("#") + + +@functools.singledispatch +def yield_lines(iterable: Iterable[str]) -> Iterator[str]: + r""" + Yield valid lines of a string or iterable. + >>> list(yield_lines('')) + [] + >>> list(yield_lines(['foo', 'bar'])) + ['foo', 'bar'] + >>> list(yield_lines('foo\nbar')) + ['foo', 'bar'] + >>> list(yield_lines('\nfoo\n#bar\nbaz #comment')) + ['foo', 'baz #comment'] + >>> list(yield_lines(['foo\nbar', 'baz', 'bing\n\n\n'])) + ['foo', 'bar', 'baz', 'bing'] + """ + return itertools.chain.from_iterable(map(yield_lines, iterable)) + + +@yield_lines.register(str) +def _(text: str) -> Iterator[str]: + return filter(_nonblank, map(str.strip, text.splitlines())) + + +def split_sections( + s: str | Iterator[str], +) -> Generator[tuple[str | None, list[str]], None, None]: + """Split a string or iterable thereof into (section, content) pairs + Each ``section`` is a stripped version of the section header ("[section]") + and each ``content`` is a list of stripped lines excluding blank lines and + comment-only lines. If there are any such lines before the first section + header, they're returned in a first ``section`` of ``None``. + """ + section = None + content: list[str] = [] + for line in yield_lines(s): + if line.startswith("["): + if line.endswith("]"): + if section or content: + yield section, content + section = line[1:-1].strip() + content = [] + else: + raise ValueError("Invalid section heading", line) + else: + content.append(line) + + # wrap up last segment + yield section, content + + +def safe_extra(extra: str) -> str: + """Convert an arbitrary string to a standard 'extra' name + Any runs of non-alphanumeric characters are replaced with a single '_', + and the result is always lowercased. + """ + return re.sub("[^A-Za-z0-9.-]+", "_", extra).lower() + + +def safe_name(name: str) -> str: + """Convert an arbitrary string to a standard distribution name + Any runs of non-alphanumeric/. characters are replaced with a single '-'. + """ + return re.sub("[^A-Za-z0-9.]+", "-", name) + + +def requires_to_requires_dist(requirement: Requirement) -> str: + """Return the version specifier for a requirement in PEP 345/566 fashion.""" + if requirement.url: + return " @ " + requirement.url + + requires_dist: list[str] = [] + for spec in requirement.specifier: + requires_dist.append(spec.operator + spec.version) + + if requires_dist: + return " " + ",".join(sorted(requires_dist)) + else: + return "" + + +def convert_requirements(requirements: list[str]) -> Iterator[str]: + """Yield Requires-Dist: strings for parsed requirements strings.""" + for req in requirements: + parsed_requirement = Requirement(req) + spec = requires_to_requires_dist(parsed_requirement) + extras = ",".join(sorted(safe_extra(e) for e in parsed_requirement.extras)) + if extras: + extras = f"[{extras}]" + + yield safe_name(parsed_requirement.name) + extras + spec + + +def generate_requirements( + extras_require: dict[str | None, list[str]], +) -> Iterator[tuple[str, str]]: + """ + Convert requirements from a setup()-style dictionary to + ('Requires-Dist', 'requirement') and ('Provides-Extra', 'extra') tuples. + + extras_require is a dictionary of {extra: [requirements]} as passed to setup(), + using the empty extra {'': [requirements]} to hold install_requires. + """ + for extra, depends in extras_require.items(): + condition = "" + extra = extra or "" + if ":" in extra: # setuptools extra:condition syntax + extra, condition = extra.split(":", 1) + + extra = safe_extra(extra) + if extra: + yield "Provides-Extra", extra + if condition: + condition = "(" + condition + ") and " + condition += f"extra == '{extra}'" + + if condition: + condition = " ; " + condition + + for new_req in convert_requirements(depends): + canonical_req = str(Requirement(new_req + condition)) + yield "Requires-Dist", canonical_req + + +def pkginfo_to_metadata(egg_info_path: str, pkginfo_path: str) -> Message: + """ + Convert .egg-info directory with PKG-INFO to the Metadata 2.1 format + """ + with open(pkginfo_path, encoding="utf-8") as headers: + pkg_info = Parser().parse(headers) + + pkg_info.replace_header("Metadata-Version", "2.1") + # Those will be regenerated from `requires.txt`. + del pkg_info["Provides-Extra"] + del pkg_info["Requires-Dist"] + requires_path = os.path.join(egg_info_path, "requires.txt") + if os.path.exists(requires_path): + with open(requires_path, encoding="utf-8") as requires_file: + requires = requires_file.read() + + parsed_requirements = sorted(split_sections(requires), key=lambda x: x[0] or "") + for extra, reqs in parsed_requirements: + for key, value in generate_requirements({extra: reqs}): + if (key, value) not in pkg_info.items(): + pkg_info[key] = value + + description = pkg_info["Description"] + if description: + description_lines = pkg_info["Description"].splitlines() + dedented_description = "\n".join( + # if the first line of long_description is blank, + # the first line here will be indented. + ( + description_lines[0].lstrip(), + textwrap.dedent("\n".join(description_lines[1:])), + "\n", + ) + ) + pkg_info.set_payload(dedented_description) + del pkg_info["Description"] + + return pkg_info diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/wheel/util.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/wheel/util.py new file mode 100644 index 0000000000000000000000000000000000000000..c928aa403b75e1ae392b4af5fe13147a095bbd30 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/wheel/util.py @@ -0,0 +1,17 @@ +from __future__ import annotations + +import base64 +import logging + +log = logging.getLogger("wheel") + + +def urlsafe_b64encode(data: bytes) -> bytes: + """urlsafe_b64encode without padding""" + return base64.urlsafe_b64encode(data).rstrip(b"=") + + +def urlsafe_b64decode(data: bytes) -> bytes: + """urlsafe_b64decode without padding""" + pad = b"=" * (4 - (len(data) & 3)) + return base64.urlsafe_b64decode(data + pad) diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/wheel/wheelfile.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/wheel/wheelfile.py new file mode 100644 index 0000000000000000000000000000000000000000..0a0f4596c566420135991713ee75ca29c4cff3ed --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/wheel/wheelfile.py @@ -0,0 +1,227 @@ +from __future__ import annotations + +import csv +import hashlib +import os.path +import re +import stat +import time +from io import StringIO, TextIOWrapper +from typing import IO, TYPE_CHECKING, Literal +from zipfile import ZIP_DEFLATED, ZipFile, ZipInfo + +from wheel.cli import WheelError +from wheel.util import log, urlsafe_b64decode, urlsafe_b64encode + +if TYPE_CHECKING: + from typing import Protocol, Sized, Union + + from typing_extensions import Buffer + + StrPath = Union[str, os.PathLike[str]] + + class SizedBuffer(Sized, Buffer, Protocol): ... + + +# Non-greedy matching of an optional build number may be too clever (more +# invalid wheel filenames will match). Separate regex for .dist-info? +WHEEL_INFO_RE = re.compile( + r"""^(?P(?P[^\s-]+?)-(?P[^\s-]+?))(-(?P\d[^\s-]*))? + -(?P[^\s-]+?)-(?P[^\s-]+?)-(?P\S+)\.whl$""", + re.VERBOSE, +) +MINIMUM_TIMESTAMP = 315532800 # 1980-01-01 00:00:00 UTC + + +def get_zipinfo_datetime(timestamp: float | None = None): + # Some applications need reproducible .whl files, but they can't do this without + # forcing the timestamp of the individual ZipInfo objects. See issue #143. + timestamp = int(os.environ.get("SOURCE_DATE_EPOCH", timestamp or time.time())) + timestamp = max(timestamp, MINIMUM_TIMESTAMP) + return time.gmtime(timestamp)[0:6] + + +class WheelFile(ZipFile): + """A ZipFile derivative class that also reads SHA-256 hashes from + .dist-info/RECORD and checks any read files against those. + """ + + _default_algorithm = hashlib.sha256 + + def __init__( + self, + file: StrPath, + mode: Literal["r", "w", "x", "a"] = "r", + compression: int = ZIP_DEFLATED, + ): + basename = os.path.basename(file) + self.parsed_filename = WHEEL_INFO_RE.match(basename) + if not basename.endswith(".whl") or self.parsed_filename is None: + raise WheelError(f"Bad wheel filename {basename!r}") + + ZipFile.__init__(self, file, mode, compression=compression, allowZip64=True) + + self.dist_info_path = "{}.dist-info".format( + self.parsed_filename.group("namever") + ) + self.record_path = self.dist_info_path + "/RECORD" + self._file_hashes: dict[str, tuple[None, None] | tuple[int, bytes]] = {} + self._file_sizes = {} + if mode == "r": + # Ignore RECORD and any embedded wheel signatures + self._file_hashes[self.record_path] = None, None + self._file_hashes[self.record_path + ".jws"] = None, None + self._file_hashes[self.record_path + ".p7s"] = None, None + + # Fill in the expected hashes by reading them from RECORD + try: + record = self.open(self.record_path) + except KeyError: + raise WheelError(f"Missing {self.record_path} file") from None + + with record: + for line in csv.reader( + TextIOWrapper(record, newline="", encoding="utf-8") + ): + path, hash_sum, size = line + if not hash_sum: + continue + + algorithm, hash_sum = hash_sum.split("=") + try: + hashlib.new(algorithm) + except ValueError: + raise WheelError( + f"Unsupported hash algorithm: {algorithm}" + ) from None + + if algorithm.lower() in {"md5", "sha1"}: + raise WheelError( + f"Weak hash algorithm ({algorithm}) is not permitted by " + f"PEP 427" + ) + + self._file_hashes[path] = ( + algorithm, + urlsafe_b64decode(hash_sum.encode("ascii")), + ) + + def open( + self, + name_or_info: str | ZipInfo, + mode: Literal["r", "w"] = "r", + pwd: bytes | None = None, + ) -> IO[bytes]: + def _update_crc(newdata: bytes) -> None: + eof = ef._eof + update_crc_orig(newdata) + running_hash.update(newdata) + if eof and running_hash.digest() != expected_hash: + raise WheelError(f"Hash mismatch for file '{ef_name}'") + + ef_name = ( + name_or_info.filename if isinstance(name_or_info, ZipInfo) else name_or_info + ) + if ( + mode == "r" + and not ef_name.endswith("/") + and ef_name not in self._file_hashes + ): + raise WheelError(f"No hash found for file '{ef_name}'") + + ef = ZipFile.open(self, name_or_info, mode, pwd) + if mode == "r" and not ef_name.endswith("/"): + algorithm, expected_hash = self._file_hashes[ef_name] + if expected_hash is not None: + # Monkey patch the _update_crc method to also check for the hash from + # RECORD + running_hash = hashlib.new(algorithm) + update_crc_orig, ef._update_crc = ef._update_crc, _update_crc + + return ef + + def write_files(self, base_dir: str): + log.info(f"creating '{self.filename}' and adding '{base_dir}' to it") + deferred: list[tuple[str, str]] = [] + for root, dirnames, filenames in os.walk(base_dir): + # Sort the directory names so that `os.walk` will walk them in a + # defined order on the next iteration. + dirnames.sort() + for name in sorted(filenames): + path = os.path.normpath(os.path.join(root, name)) + if os.path.isfile(path): + arcname = os.path.relpath(path, base_dir).replace(os.path.sep, "/") + if arcname == self.record_path: + pass + elif root.endswith(".dist-info"): + deferred.append((path, arcname)) + else: + self.write(path, arcname) + + deferred.sort() + for path, arcname in deferred: + self.write(path, arcname) + + def write( + self, + filename: str, + arcname: str | None = None, + compress_type: int | None = None, + ) -> None: + with open(filename, "rb") as f: + st = os.fstat(f.fileno()) + data = f.read() + + zinfo = ZipInfo( + arcname or filename, date_time=get_zipinfo_datetime(st.st_mtime) + ) + zinfo.external_attr = (stat.S_IMODE(st.st_mode) | stat.S_IFMT(st.st_mode)) << 16 + zinfo.compress_type = compress_type or self.compression + self.writestr(zinfo, data, compress_type) + + def writestr( + self, + zinfo_or_arcname: str | ZipInfo, + data: SizedBuffer | str, + compress_type: int | None = None, + ): + if isinstance(zinfo_or_arcname, str): + zinfo_or_arcname = ZipInfo( + zinfo_or_arcname, date_time=get_zipinfo_datetime() + ) + zinfo_or_arcname.compress_type = self.compression + zinfo_or_arcname.external_attr = (0o664 | stat.S_IFREG) << 16 + + if isinstance(data, str): + data = data.encode("utf-8") + + ZipFile.writestr(self, zinfo_or_arcname, data, compress_type) + fname = ( + zinfo_or_arcname.filename + if isinstance(zinfo_or_arcname, ZipInfo) + else zinfo_or_arcname + ) + log.info(f"adding '{fname}'") + if fname != self.record_path: + hash_ = self._default_algorithm(data) + self._file_hashes[fname] = ( + hash_.name, + urlsafe_b64encode(hash_.digest()).decode("ascii"), + ) + self._file_sizes[fname] = len(data) + + def close(self): + # Write RECORD + if self.fp is not None and self.mode == "w" and self._file_hashes: + data = StringIO() + writer = csv.writer(data, delimiter=",", quotechar='"', lineterminator="\n") + writer.writerows( + ( + (fname, algorithm + "=" + hash_, self._file_sizes[fname]) + for fname, (algorithm, hash_) in self._file_hashes.items() + ) + ) + writer.writerow((format(self.record_path), "", "")) + self.writestr(self.record_path, data.getvalue()) + + ZipFile.close(self) diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/yarl/_parse.py b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/yarl/_parse.py new file mode 100644 index 0000000000000000000000000000000000000000..cc259ea86dc89aa0766d0cd30945fd605743b7b0 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/yarl/_parse.py @@ -0,0 +1,189 @@ +"""URL parsing utilities.""" + +import re +import unicodedata +from functools import lru_cache +from typing import Union +from urllib.parse import scheme_chars, uses_netloc + +from ._quoters import QUOTER + +# Leading and trailing C0 control and space to be stripped per WHATWG spec. +# == "".join([chr(i) for i in range(0, 0x20 + 1)]) +WHATWG_C0_CONTROL_OR_SPACE = ( + "\x00\x01\x02\x03\x04\x05\x06\x07\x08\t\n\x0b\x0c\r\x0e\x0f\x10" + "\x11\x12\x13\x14\x15\x16\x17\x18\x19\x1a\x1b\x1c\x1d\x1e\x1f " +) + +# Unsafe bytes to be removed per WHATWG spec +UNSAFE_URL_BYTES_TO_REMOVE = ["\t", "\r", "\n"] +USES_AUTHORITY = frozenset(uses_netloc) + +SplitURLType = tuple[str, str, str, str, str] + + +def split_url(url: str) -> SplitURLType: + """Split URL into parts.""" + # Adapted from urllib.parse.urlsplit + # Only lstrip url as some applications rely on preserving trailing space. + # (https://url.spec.whatwg.org/#concept-basic-url-parser would strip both) + url = url.lstrip(WHATWG_C0_CONTROL_OR_SPACE) + for b in UNSAFE_URL_BYTES_TO_REMOVE: + if b in url: + url = url.replace(b, "") + + scheme = netloc = query = fragment = "" + i = url.find(":") + if i > 0 and url[0] in scheme_chars: + for c in url[1:i]: + if c not in scheme_chars: + break + else: + scheme, url = url[:i].lower(), url[i + 1 :] + has_hash = "#" in url + has_question_mark = "?" in url + if url[:2] == "//": + delim = len(url) # position of end of domain part of url, default is end + if has_hash and has_question_mark: + delim_chars = "/?#" + elif has_question_mark: + delim_chars = "/?" + elif has_hash: + delim_chars = "/#" + else: + delim_chars = "/" + for c in delim_chars: # look for delimiters; the order is NOT important + wdelim = url.find(c, 2) # find first of this delim + if wdelim >= 0 and wdelim < delim: # if found + delim = wdelim # use earliest delim position + netloc = url[2:delim] + url = url[delim:] + has_left_bracket = "[" in netloc + has_right_bracket = "]" in netloc + if (has_left_bracket and not has_right_bracket) or ( + has_right_bracket and not has_left_bracket + ): + raise ValueError("Invalid IPv6 URL") + if has_left_bracket: + bracketed_host = netloc.partition("[")[2].partition("]")[0] + # Valid bracketed hosts are defined in + # https://www.rfc-editor.org/rfc/rfc3986#page-49 + # https://url.spec.whatwg.org/ + if bracketed_host[0] == "v": + if not re.match(r"\Av[a-fA-F0-9]+\..+\Z", bracketed_host): + raise ValueError("IPvFuture address is invalid") + elif ":" not in bracketed_host: + raise ValueError("An IPv4 address cannot be in brackets") + if has_hash: + url, _, fragment = url.partition("#") + if has_question_mark: + url, _, query = url.partition("?") + if netloc and not netloc.isascii(): + _check_netloc(netloc) + return scheme, netloc, url, query, fragment + + +def _check_netloc(netloc: str) -> None: + # Adapted from urllib.parse._checknetloc + # looking for characters like \u2100 that expand to 'a/c' + # IDNA uses NFKC equivalence, so normalize for this check + + # ignore characters already included + # but not the surrounding text + n = netloc.replace("@", "").replace(":", "").replace("#", "").replace("?", "") + normalized_netloc = unicodedata.normalize("NFKC", n) + if n == normalized_netloc: + return + # Note that there are no unicode decompositions for the character '@' so + # its currently impossible to have test coverage for this branch, however if the + # one should be added in the future we want to make sure its still checked. + for c in "/?#@:": # pragma: no branch + if c in normalized_netloc: + raise ValueError( + f"netloc '{netloc}' contains invalid " + "characters under NFKC normalization" + ) + + +@lru_cache # match the same size as urlsplit +def split_netloc( + netloc: str, +) -> tuple[Union[str, None], Union[str, None], Union[str, None], Union[int, None]]: + """Split netloc into username, password, host and port.""" + if "@" not in netloc: + username: Union[str, None] = None + password: Union[str, None] = None + hostinfo = netloc + else: + userinfo, _, hostinfo = netloc.rpartition("@") + username, have_password, password = userinfo.partition(":") + if not have_password: + password = None + + if "[" in hostinfo: + _, _, bracketed = hostinfo.partition("[") + hostname, _, port_str = bracketed.partition("]") + _, _, port_str = port_str.partition(":") + else: + hostname, _, port_str = hostinfo.partition(":") + + if not port_str: + return username or None, password, hostname or None, None + + try: + port = int(port_str) + except ValueError: + raise ValueError("Invalid URL: port can't be converted to integer") + if not (0 <= port <= 65535): + raise ValueError("Port out of range 0-65535") + return username or None, password, hostname or None, port + + +def unsplit_result( + scheme: str, netloc: str, url: str, query: str, fragment: str +) -> str: + """Unsplit a URL without any normalization.""" + if netloc or (scheme and scheme in USES_AUTHORITY) or url[:2] == "//": + if url and url[:1] != "/": + url = f"{scheme}://{netloc}/{url}" if scheme else f"{scheme}:{url}" + else: + url = f"{scheme}://{netloc}{url}" if scheme else f"//{netloc}{url}" + elif scheme: + url = f"{scheme}:{url}" + if query: + url = f"{url}?{query}" + return f"{url}#{fragment}" if fragment else url + + +@lru_cache # match the same size as urlsplit +def make_netloc( + user: Union[str, None], + password: Union[str, None], + host: Union[str, None], + port: Union[int, None], + encode: bool = False, +) -> str: + """Make netloc from parts. + + The user and password are encoded if encode is True. + + The host must already be encoded with _encode_host. + """ + if host is None: + return "" + ret = host + if port is not None: + ret = f"{ret}:{port}" + if user is None and password is None: + return ret + if password is not None: + if not user: + user = "" + elif encode: + user = QUOTER(user) + if encode: + password = QUOTER(password) + user = f"{user}:{password}" + elif user and encode: + user = QUOTER(user) + return f"{user}@{ret}" if user else ret diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/yarl/_quoting_c.pyx b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/yarl/_quoting_c.pyx new file mode 100644 index 0000000000000000000000000000000000000000..067ba96e4efa297669b22fa407af500c7ea1e636 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/yarl/_quoting_c.pyx @@ -0,0 +1,423 @@ +# cython: language_level=3 + +from cpython.exc cimport PyErr_NoMemory +from cpython.mem cimport PyMem_Free, PyMem_Malloc, PyMem_Realloc +from cpython.unicode cimport ( + PyUnicode_DATA, + PyUnicode_DecodeASCII, + PyUnicode_DecodeUTF8Stateful, + PyUnicode_GET_LENGTH, + PyUnicode_KIND, + PyUnicode_READ, +) +from libc.stdint cimport uint8_t, uint64_t +from libc.string cimport memcpy, memset + +from string import ascii_letters, digits + + +cdef str GEN_DELIMS = ":/?#[]@" +cdef str SUB_DELIMS_WITHOUT_QS = "!$'()*," +cdef str SUB_DELIMS = SUB_DELIMS_WITHOUT_QS + '+?=;' +cdef str RESERVED = GEN_DELIMS + SUB_DELIMS +cdef str UNRESERVED = ascii_letters + digits + '-._~' +cdef str ALLOWED = UNRESERVED + SUB_DELIMS_WITHOUT_QS +cdef str QS = '+&=;' + +DEF BUF_SIZE = 8 * 1024 # 8KiB +cdef char BUFFER[BUF_SIZE] + +cdef inline Py_UCS4 _to_hex(uint8_t v) noexcept: + if v < 10: + return (v+0x30) # ord('0') == 0x30 + else: + return (v+0x41-10) # ord('A') == 0x41 + + +cdef inline int _from_hex(Py_UCS4 v) noexcept: + if '0' <= v <= '9': + return (v) - 0x30 # ord('0') == 0x30 + elif 'A' <= v <= 'F': + return (v) - 0x41 + 10 # ord('A') == 0x41 + elif 'a' <= v <= 'f': + return (v) - 0x61 + 10 # ord('a') == 0x61 + else: + return -1 + + +cdef inline int _is_lower_hex(Py_UCS4 v) noexcept: + return 'a' <= v <= 'f' + + +cdef inline Py_UCS4 _restore_ch(Py_UCS4 d1, Py_UCS4 d2): + cdef int digit1 = _from_hex(d1) + if digit1 < 0: + return -1 + cdef int digit2 = _from_hex(d2) + if digit2 < 0: + return -1 + return (digit1 << 4 | digit2) + + +cdef uint8_t ALLOWED_TABLE[16] +cdef uint8_t ALLOWED_NOTQS_TABLE[16] + + +cdef inline bint bit_at(uint8_t array[], uint64_t ch) noexcept: + return array[ch >> 3] & (1 << (ch & 7)) + + +cdef inline void set_bit(uint8_t array[], uint64_t ch) noexcept: + array[ch >> 3] |= (1 << (ch & 7)) + + +memset(ALLOWED_TABLE, 0, sizeof(ALLOWED_TABLE)) +memset(ALLOWED_NOTQS_TABLE, 0, sizeof(ALLOWED_NOTQS_TABLE)) + +for i in range(128): + if chr(i) in ALLOWED: + set_bit(ALLOWED_TABLE, i) + set_bit(ALLOWED_NOTQS_TABLE, i) + if chr(i) in QS: + set_bit(ALLOWED_NOTQS_TABLE, i) + +# ----------------- writer --------------------------- + +cdef struct Writer: + char *buf + Py_ssize_t size + Py_ssize_t pos + bint changed + + +cdef inline void _init_writer(Writer* writer): + writer.buf = &BUFFER[0] + writer.size = BUF_SIZE + writer.pos = 0 + writer.changed = 0 + + +cdef inline void _release_writer(Writer* writer): + if writer.buf != BUFFER: + PyMem_Free(writer.buf) + + +cdef inline int _write_char(Writer* writer, Py_UCS4 ch, bint changed): + cdef char * buf + cdef Py_ssize_t size + + if writer.pos == writer.size: + # reallocate + size = writer.size + BUF_SIZE + if writer.buf == BUFFER: + buf = PyMem_Malloc(size) + if buf == NULL: + PyErr_NoMemory() + return -1 + memcpy(buf, writer.buf, writer.size) + else: + buf = PyMem_Realloc(writer.buf, size) + if buf == NULL: + PyErr_NoMemory() + return -1 + writer.buf = buf + writer.size = size + writer.buf[writer.pos] = ch + writer.pos += 1 + writer.changed |= changed + return 0 + + +cdef inline int _write_pct(Writer* writer, uint8_t ch, bint changed): + if _write_char(writer, '%', changed) < 0: + return -1 + if _write_char(writer, _to_hex(ch >> 4), changed) < 0: + return -1 + return _write_char(writer, _to_hex(ch & 0x0f), changed) + + +cdef inline int _write_utf8(Writer* writer, Py_UCS4 symbol): + cdef uint64_t utf = symbol + + if utf < 0x80: + return _write_pct(writer, utf, True) + elif utf < 0x800: + if _write_pct(writer, (0xc0 | (utf >> 6)), True) < 0: + return -1 + return _write_pct(writer, (0x80 | (utf & 0x3f)), True) + elif 0xD800 <= utf <= 0xDFFF: + # surogate pair, ignored + return 0 + elif utf < 0x10000: + if _write_pct(writer, (0xe0 | (utf >> 12)), True) < 0: + return -1 + if _write_pct(writer, (0x80 | ((utf >> 6) & 0x3f)), + True) < 0: + return -1 + return _write_pct(writer, (0x80 | (utf & 0x3f)), True) + elif utf > 0x10FFFF: + # symbol is too large + return 0 + else: + if _write_pct(writer, (0xf0 | (utf >> 18)), True) < 0: + return -1 + if _write_pct(writer, (0x80 | ((utf >> 12) & 0x3f)), + True) < 0: + return -1 + if _write_pct(writer, (0x80 | ((utf >> 6) & 0x3f)), + True) < 0: + return -1 + return _write_pct(writer, (0x80 | (utf & 0x3f)), True) + + +# --------------------- end writer -------------------------- + + +cdef class _Quoter: + cdef bint _qs + cdef bint _requote + + cdef uint8_t _safe_table[16] + cdef uint8_t _protected_table[16] + + def __init__( + self, *, str safe='', str protected='', bint qs=False, bint requote=True, + ): + cdef Py_UCS4 ch + + self._qs = qs + self._requote = requote + + if not self._qs: + memcpy(self._safe_table, + ALLOWED_NOTQS_TABLE, + sizeof(self._safe_table)) + else: + memcpy(self._safe_table, + ALLOWED_TABLE, + sizeof(self._safe_table)) + for ch in safe: + if ord(ch) > 127: + raise ValueError("Only safe symbols with ORD < 128 are allowed") + set_bit(self._safe_table, ch) + + memset(self._protected_table, 0, sizeof(self._protected_table)) + for ch in protected: + if ord(ch) > 127: + raise ValueError("Only safe symbols with ORD < 128 are allowed") + set_bit(self._safe_table, ch) + set_bit(self._protected_table, ch) + + def __call__(self, val): + if val is None: + return None + if type(val) is not str: + if isinstance(val, str): + # derived from str + val = str(val) + else: + raise TypeError("Argument should be str") + return self._do_quote_or_skip(val) + + cdef str _do_quote_or_skip(self, str val): + cdef Py_UCS4 ch + cdef Py_ssize_t length = PyUnicode_GET_LENGTH(val) + cdef Py_ssize_t idx = length + cdef bint must_quote = 0 + cdef Writer writer + cdef int kind = PyUnicode_KIND(val) + cdef const void *data = PyUnicode_DATA(val) + + # If everything in the string is in the safe + # table and all ASCII, we can skip quoting + while idx: + idx -= 1 + ch = PyUnicode_READ(kind, data, idx) + if ch >= 128 or not bit_at(self._safe_table, ch): + must_quote = 1 + break + + if not must_quote: + return val + + _init_writer(&writer) + try: + return self._do_quote(val, length, kind, data, &writer) + finally: + _release_writer(&writer) + + cdef str _do_quote( + self, + str val, + Py_ssize_t length, + int kind, + const void *data, + Writer *writer + ): + cdef Py_UCS4 ch + cdef int changed + cdef Py_ssize_t idx = 0 + + while idx < length: + ch = PyUnicode_READ(kind, data, idx) + idx += 1 + if ch == '%' and self._requote and idx <= length - 2: + ch = _restore_ch( + PyUnicode_READ(kind, data, idx), + PyUnicode_READ(kind, data, idx + 1) + ) + if ch != -1: + idx += 2 + if ch < 128: + if bit_at(self._protected_table, ch): + if _write_pct(writer, ch, True) < 0: + raise + continue + + if bit_at(self._safe_table, ch): + if _write_char(writer, ch, True) < 0: + raise + continue + + changed = (_is_lower_hex(PyUnicode_READ(kind, data, idx - 2)) or + _is_lower_hex(PyUnicode_READ(kind, data, idx - 1))) + if _write_pct(writer, ch, changed) < 0: + raise + continue + else: + ch = '%' + + if self._write(writer, ch) < 0: + raise + + if not writer.changed: + return val + else: + return PyUnicode_DecodeASCII(writer.buf, writer.pos, "strict") + + cdef inline int _write(self, Writer *writer, Py_UCS4 ch): + if self._qs: + if ch == ' ': + return _write_char(writer, '+', True) + + if ch < 128 and bit_at(self._safe_table, ch): + return _write_char(writer, ch, False) + + return _write_utf8(writer, ch) + + +cdef class _Unquoter: + cdef str _ignore + cdef str _unsafe + cdef bint _qs + cdef _Quoter _quoter + cdef _Quoter _qs_quoter + + def __init__(self, *, ignore="", unsafe="", qs=False): + self._ignore = ignore + self._unsafe = unsafe + self._qs = qs + self._quoter = _Quoter() + self._qs_quoter = _Quoter(qs=True) + + def __call__(self, val): + if val is None: + return None + if type(val) is not str: + if isinstance(val, str): + # derived from str + val = str(val) + else: + raise TypeError("Argument should be str") + return self._do_unquote(val) + + cdef str _do_unquote(self, str val): + cdef Py_ssize_t length = PyUnicode_GET_LENGTH(val) + if length == 0: + return val + + cdef list ret = [] + cdef char buffer[4] + cdef Py_ssize_t buflen = 0 + cdef Py_ssize_t consumed + cdef str unquoted + cdef Py_UCS4 ch = 0 + cdef Py_ssize_t idx = 0 + cdef Py_ssize_t start_pct + cdef int kind = PyUnicode_KIND(val) + cdef const void *data = PyUnicode_DATA(val) + cdef bint changed = 0 + while idx < length: + ch = PyUnicode_READ(kind, data, idx) + idx += 1 + if ch == '%' and idx <= length - 2: + changed = 1 + ch = _restore_ch( + PyUnicode_READ(kind, data, idx), + PyUnicode_READ(kind, data, idx + 1) + ) + if ch != -1: + idx += 2 + assert buflen < 4 + buffer[buflen] = ch + buflen += 1 + try: + unquoted = PyUnicode_DecodeUTF8Stateful(buffer, buflen, + NULL, &consumed) + except UnicodeDecodeError: + start_pct = idx - buflen * 3 + buffer[0] = ch + buflen = 1 + ret.append(val[start_pct : idx - 3]) + try: + unquoted = PyUnicode_DecodeUTF8Stateful(buffer, buflen, + NULL, &consumed) + except UnicodeDecodeError: + buflen = 0 + ret.append(val[idx - 3 : idx]) + continue + if not unquoted: + assert consumed == 0 + continue + assert consumed == buflen + buflen = 0 + if self._qs and unquoted in '+=&;': + ret.append(self._qs_quoter(unquoted)) + elif unquoted in self._unsafe or unquoted in self._ignore: + ret.append(self._quoter(unquoted)) + else: + ret.append(unquoted) + continue + else: + ch = '%' + + if buflen: + start_pct = idx - 1 - buflen * 3 + ret.append(val[start_pct : idx - 1]) + buflen = 0 + + if ch == '+': + if not self._qs or ch in self._unsafe: + ret.append('+') + else: + changed = 1 + ret.append(' ') + continue + + if ch in self._unsafe: + changed = 1 + ret.append('%') + h = hex(ord(ch)).upper()[2:] + for ch in h: + ret.append(ch) + continue + + ret.append(ch) + + if not changed: + return val + + if buflen: + ret.append(val[length - buflen * 3 : length]) + + return ''.join(ret) diff --git a/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/yarl/py.typed b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/yarl/py.typed new file mode 100644 index 0000000000000000000000000000000000000000..dcf2c804da5e19d617a03a6c68aa128d1d1f89a0 --- /dev/null +++ b/Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/yarl/py.typed @@ -0,0 +1 @@ +# Placeholder