科研工程能力：集群任务调度(Slurm)

import subprocess # 模拟通过 sbatch 提交一个 GPU 分布式训练作业脚本script = """#!/bin/bash#SBATCH --job-name=train_gpt#SBATCH --partition=gpu#SBATCH --nodes=4#SBATCH --ntasks-per-node=4#SBATCH --cpus-per-task=8#SBATCH --gres=gpu:4#SBATCH --mem=128G#SBATCH --time=24:00:00#SBATCH --output=train_%j.out module load cuda/12.1torchrun --nproc_per_node=4 train.py"""with open("train_job.slurm", "w") as f:    f.write(script) result = subprocess.run(["sbatch", "train_job.slurm"], capture_output=True, text=True)print(result.stdout.strip())

import subprocessimport timefrom dataclasses import dataclassfrom typing import Dict  @dataclassclass Job:    job_id: int    name: str    partition: str    num_nodes: int    num_tasks: int    cpus_per_task: int    gpus_per_node: int    mem_per_node: str    time_limit: str    priority: float = 0.0    state: str = "PENDING"  class SlurmScheduler:    def __init__(self):        self.jobs: Dict[int, Job] = {}        self.partitions: Dict[str, Dict] = {            "gpu": {"max_nodes": 8, "gpus_per_node": 8, "max_time": "72:00:00"},            "cpu": {"max_nodes": 16, "gpus_per_node": 0, "max_time": "168:00:00"},        }        self.next_job_id = 1        self.fairshare_weights = {"gpu": 0.5, "cpu": 0.7}     def submit_job(self, name, partition, num_nodes, num_tasks,                   cpus_per_task, gpus_per_node, mem_per_node, time_limit):        job = Job(            job_id=self.next_job_id,            name=name,            partition=partition,            num_nodes=num_nodes,            num_tasks=num_tasks,            cpus_per_task=cpus_per_task,            gpus_per_node=gpus_per_node,            mem_per_node=mem_per_node,            time_limit=time_limit,        )        self.next_job_id += 1        self.jobs[job.job_id] = job        self._compute_priority(job)        return job.job_id     def _compute_priority(self, job):        age_factor = 0.3        fairshare_factor = 0.5        job_size_factor = 0.2        age = (self.next_job_id - job.job_id) / max(self.next_job_id, 1)        fairshare = self.fairshare_weights.get(job.partition, 0.5)        size = job.num_nodes * job.cpus_per_task / 100.0        job.priority = (age_factor * age + fairshare_factor * fairshare                        + job_size_factor * size)     def allocate_resources(self, job):        partition = self.partitions.get(job.partition)        if partition is None:            return False        if job.num_nodes > partition["max_nodes"]:            return False        if job.gpus_per_node > partition["gpus_per_node"]:            return False        job.state = "RUNNING"        return True     def run_scheduler(self):        pending = [j for j in self.jobs.values() if j.state == "PENDING"]        pending.sort(key=lambda j: j.priority, reverse=True)        for job in pending:            if self.allocate_resources(job):                print(f"Job {job.job_id} ({job.name}) -> RUNNING "                      f"on {job.partition}")     def list_jobs(self):        print("ID\tName\t\tPartition\tState\t\tPriority")        for job in self.jobs.values():            print(f"{job.job_id}\t{job.name}\t{job.partition}\t\t"                  f"{job.state}\t{job.priority:.3f}")  if __name__ == "__main__":    scheduler = SlurmScheduler()    scheduler.submit_job("train_gpt", "gpu", 4, 4, 8, 4, "128G", "24:00:00")    scheduler.submit_job("eval_model", "gpu", 1, 1, 8, 1, "64G", "02:00:00")    scheduler.submit_job("data_proc", "cpu", 2, 8, 4, 0, "32G", "12:00:00")    scheduler.run_scheduler()    scheduler.list_jobs()