lerobot/examples/port_datasets/openx_utils/test.py

55 lines
1.5 KiB
Python

import os
import psutil
def display_system_info():
# Get the number of CPUs
num_cpus = psutil.cpu_count(logical=True)
print(f"Number of CPUs: {num_cpus}")
# Get memory information
memory_info = psutil.virtual_memory()
total_memory = memory_info.total / (1024**3) # Convert bytes to GB
available_memory = memory_info.available / (1024**3) # Convert bytes to GB
used_memory = memory_info.used / (1024**3) # Convert bytes to GB
print(f"Total Memory: {total_memory:.2f} GB")
print(f"Available Memory: {available_memory:.2f} GB")
print(f"Used Memory: {used_memory:.2f} GB")
def display_slurm_info():
# Get SLURM job ID
job_id = os.getenv("SLURM_JOB_ID")
print(f"SLURM Job ID: {job_id}")
# Get SLURM job name
job_name = os.getenv("SLURM_JOB_NAME")
print(f"SLURM Job Name: {job_name}")
# Get the number of tasks
num_tasks = os.getenv("SLURM_NTASKS")
print(f"Number of Tasks: {num_tasks}")
# Get the number of nodes
num_nodes = os.getenv("SLURM_NNODES")
print(f"Number of Nodes: {num_nodes}")
# Get the number of CPUs per task
cpus_per_task = os.getenv("SLURM_CPUS_PER_TASK")
print(f"CPUs per Task: {cpus_per_task}")
# Get the node list
node_list = os.getenv("SLURM_NODELIST")
print(f"Node List: {node_list}")
# Get the task ID (only available within an srun task)
task_id = os.getenv("SLURM_PROCID")
print(f"Task ID: {task_id}")
if __name__ == "__main__":
display_system_info()
display_slurm_info()