# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
from gcm.exporters.graph_api import GraphAPI
from gcm.monitoring.clock import ClockImpl
from gcm.monitoring.sink.protocol import DataType, SinkAdditionalParams
from gcm.schemas.log import Log
from gcm.schemas.slurm.sacct import SacctPayload
from gcm.tests.config import Config
from gcm.tests.conftest import report_url

SCRIBE_CATEGORY = "perfpipe_fair_cluster_sacct_test"
TEST_TIME = ClockImpl().unixtime()
SACCT_DATA = Log(
    ts=TEST_TIME,
    message=[
        SacctPayload(
            time=TEST_TIME,
            end_ds=str(ClockImpl().unixtime()),
            cluster="test",
            sacct={
                "Account": "fairusers",
                "AdminComment": "",
                "AllocCPUS": "10",
                "AllocNodes": "1",
                "AllocTRES": "cpu=10,gres/gpu=1,mem=0,node=1",
                "AssocID": "210",
                "AveCPU": "00:00:00",
                "AveCPUFreq": "2.37M",
                "AveDiskRead": "0.02M",
                "AveDiskWrite": "0.00M",
                "AvePages": "0",
                "AveRSS": "5428K",
                "AveVMSize": "6220K",
                "BlockID": "",
                "CPUTime": "15-00:04:40",
                "CPUTimeRAW": "1296280",
                "Cluster": "h2node",
                "Comment": "",
                "Constraints": "",
                "ConsumedEnergy": "0",
                "ConsumedEnergyRaw": "0",
                "Container": "",
                "DBIndex": "479788995",
                "DerivedExitCode": "",
                "Elapsed": "1-12:00:28",
                "ElapsedRaw": "129628",
                "Eligible": "2023-04-29T20:52:02-07:00",
                "End": "2023-05-01T08:52:30-07:00",
                "ExitCode": "0:15",
                "Flags": "",
                "GID": "",
                "Group": "",
                "JobID": "7195895_6.batch",
                "JobIDRaw": "7196062.batch",
                "JobName": "batch",
                "Layout": "Unknown",
                "MaxDiskRead": "0.02M",
                "MaxDiskReadNode": "node0584",
                "MaxDiskReadTask": "0",
                "MaxDiskWrite": "0.00M",
                "MaxDiskWriteNode": "node0584",
                "MaxDiskWriteTask": "0",
                "MaxPages": "0",
                "MaxPagesNode": "node0584",
                "MaxPagesTask": "0",
                "MaxRSS": "5428K",
                "MaxRSSNode": "node0584",
                "MaxRSSTask": "0",
                "MaxVMSize": "6220K",
                "MaxVMSizeNode": "node0584",
                "MaxVMSizeTask": "0",
                "McsLabel": "",
                "MinCPU": "00:00:00",
                "MinCPUNode": "node0584",
                "MinCPUTask": "0",
                "NCPUS": "10",
                "NNodes": "1",
                "NTasks": "1",
                "NodeList": "node0584",
                "Partition": "",
                "Priority": "",
                "QOS": "",
                "QOSRAW": "",
                "Reason": "",
                "ReqCPUFreq": "0",
                "ReqCPUFreqGov": "0",
                "ReqCPUFreqMax": "0",
                "ReqCPUFreqMin": "0",
                "ReqCPUS": "10",
                "ReqMem": "",
                "ReqNodes": "1",
                "ReqTRES": "",
                "Reservation": "",
                "ReservationId": "",
                "Reserved": "",
                "ResvCPU": "",
                "ResvCPURAW": "",
                "Start": "2023-04-29T20:52:02-07:00",
                "State": "CANCELLED",
                "Submit": "2023-04-29T20:52:02-07:00",
                "SubmitLine": "",
                "Suspended": "00:00:00",
                "SystemCPU": "00:00.002",
                "SystemComment": "",
                "TRESUsageInAve": "cpu=00:00:00,energy=0,fs/disk=18723,mem=5428K,pages=0,vmem=6220K",
                "TRESUsageInMax": "cpu=00:00:00,energy=0,fs/disk=18723,mem=5428K,pages=0,vmem=6220K",
                "TRESUsageInMaxNode": "cpu=node0584,energy=node0584,fs/disk=node0584,mem=node0584,pages=node0584,vmem=node0584",
                "TRESUsageInMaxTask": "cpu=0,fs/disk=0,mem=0,pages=0,vmem=0",
                "TRESUsageInMin": "cpu=00:00:00,energy=0,fs/disk=18723,mem=5428K,pages=0,vmem=6220K",
                "TRESUsageInMinNode": "cpu=node0584,energy=node0584,fs/disk=node0584,mem=node0584,pages=node0584,vmem=node0584",
                "TRESUsageInMinTask": "cpu=0,fs/disk=0,mem=0,pages=0,vmem=0",
                "TRESUsageInTot": "cpu=00:00:00,energy=0,fs/disk=18723,mem=5428K,pages=0,vmem=6220K",
                "TRESUsageOutAve": "energy=0,fs/disk=1",
                "TRESUsageOutMax": "energy=0,fs/disk=1",
                "TRESUsageOutMaxNode": "energy=node0584,fs/disk=node0584",
                "TRESUsageOutMaxTask": "fs/disk=0",
                "TRESUsageOutMin": "energy=0,fs/disk=1",
                "TRESUsageOutMinNode": "energy=node0584,fs/disk=node0584",
                "TRESUsageOutMinTask": "fs/disk=0",
                "TRESUsageOutTot": "energy=0,fs/disk=1",
                "Timelimit": "",
                "TimelimitRaw": "",
                "TotalCPU": "00:00.005",
                "UID": "",
                "User": "",
                "UserCPU": "00:00.002",
                "WCKey": "",
                "WCKeyID": "",
                "WorkDir": "",
            },
        ),
        SacctPayload(
            time=TEST_TIME,
            end_ds=str(ClockImpl().unixtime()),
            cluster="test",
            sacct={
                "Account": "fairusers",
                "AdminComment": "",
                "AllocCPUS": "10",
                "AllocNodes": "1",
                "AllocTRES": "cpu=10,gres/gpu=1,mem=0,node=1",
                "AssocID": "210",
                "AveCPU": "1-08:21:24",
                "AveCPUFreq": "9K",
                "AveDiskRead": "4.56M",
                "AveDiskWrite": "0.02M",
                "AvePages": "2145",
                "AveRSS": "3313748K",
                "AveVMSize": "35994228K",
                "BlockID": "",
                "CPUTime": "15-00:09:00",
                "CPUTimeRAW": "1296540",
                "Cluster": "h2node",
                "Comment": "",
                "Constraints": "",
                "ConsumedEnergy": "0",
                "ConsumedEnergyRaw": "0",
                "Container": "",
                "DBIndex": "479788995",
                "DerivedExitCode": "",
                "Elapsed": "1-12:00:54",
                "ElapsedRaw": "129654",
                "Eligible": "2023-04-29T20:52:15-07:00",
                "End": "2023-05-01T08:53:09-07:00",
                "ExitCode": "0:9",
                "Flags": "",
                "GID": "",
                "Group": "",
                "JobID": "7195895_6.0",
                "JobIDRaw": "7196062.0",
                "JobName": "driver_mvp.sh",
                "Layout": "Cyclic",
                "MaxDiskRead": "4.56M",
                "MaxDiskReadNode": "node0584",
                "MaxDiskReadTask": "0",
                "MaxDiskWrite": "0.02M",
                "MaxDiskWriteNode": "node0584",
                "MaxDiskWriteTask": "0",
                "MaxPages": "2145",
                "MaxPagesNode": "node0584",
                "MaxPagesTask": "0",
                "MaxRSS": "3313748K",
                "MaxRSSNode": "node0584",
                "MaxRSSTask": "0",
                "MaxVMSize": "35994228K",
                "MaxVMSizeNode": "node0584",
                "MaxVMSizeTask": "0",
                "McsLabel": "",
                "MinCPU": "1-08:21:24",
                "MinCPUNode": "node0584",
                "MinCPUTask": "0",
                "NCPUS": "10",
                "NNodes": "1",
                "NTasks": "1",
                "NodeList": "node0584",
                "Partition": "",
                "Priority": "",
                "QOS": "",
                "QOSRAW": "",
                "Reason": "",
                "ReqCPUFreq": "Unknown",
                "ReqCPUFreqGov": "Unknown",
                "ReqCPUFreqMax": "Unknown",
                "ReqCPUFreqMin": "Unknown",
                "ReqCPUS": "10",
                "ReqMem": "",
                "ReqNodes": "1",
                "ReqTRES": "",
                "Reservation": "",
                "ReservationId": "",
                "Reserved": "",
                "ResvCPU": "",
                "ResvCPURAW": "",
                "Start": "None",
                "State": "CANCELLED",
                "Submit": "2023-04-29T20:52:15-07:00",
                "SubmitLine": "srun --label /private/home/test_user/devfair/driver_mvp.sh 1 pixel FrankaMovePixels model_zoo/mae_vitl/PixMC3NormInit/FrankaMove/2e-3 train.learn.optim_stepsize=2e-3 train.encoder.size=vitl train.encoder.pretrain_path=/checkpoint/test_user/video/model_zoo/mae_vitl/pretrained.pth train.encoder.load_norm=false",
                "Suspended": "00:00:00",
                "SystemCPU": "03:23.061",
                "SystemComment": "",
                "TRESUsageInAve": "cpu=1-08:21:24,energy=0,fs/disk=4778758,mem=3313748K,pages=2145,vmem=35994228K",
                "TRESUsageInMax": "cpu=1-08:21:24,energy=0,fs/disk=4778758,mem=3313748K,pages=2145,vmem=35994228K",
                "TRESUsageInMaxNode": "cpu=node0584,energy=node0584,fs/disk=node0584,mem=node0584,pages=node0584,vmem=node0584",
                "TRESUsageInMaxTask": "cpu=0,fs/disk=0,mem=0,pages=0,vmem=0",
                "TRESUsageInMin": "cpu=1-08:21:24,energy=0,fs/disk=4778758,mem=3313748K,pages=2145,vmem=35994228K",
                "TRESUsageInMinNode": "cpu=node0584,energy=node0584,fs/disk=node0584,mem=node0584,pages=node0584,vmem=node0584",
                "TRESUsageInMinTask": "cpu=0,fs/disk=0,mem=0,pages=0,vmem=0",
                "TRESUsageInTot": "cpu=1-08:21:24,energy=0,fs/disk=4778758,mem=3313748K,pages=2145,vmem=35994228K",
                "TRESUsageOutAve": "energy=0,fs/disk=19010",
                "TRESUsageOutMax": "energy=0,fs/disk=19010",
                "TRESUsageOutMaxNode": "energy=node0584,fs/disk=node0584",
                "TRESUsageOutMaxTask": "fs/disk=0",
                "TRESUsageOutMin": "energy=0,fs/disk=19010",
                "TRESUsageOutMinNode": "energy=node0584,fs/disk=node0584",
                "TRESUsageOutMinTask": "fs/disk=0",
                "TRESUsageOutTot": "energy=0,fs/disk=19010",
                "Timelimit": "",
                "TimelimitRaw": "",
                "TotalCPU": "1-08:21:24",
                "UID": "",
                "User": "",
                "UserCPU": "1-08:18:01",
                "WCKey": "",
                "WCKeyID": "",
                "WorkDir": "",
            },
        ),
    ],
)


class TestSacctE2e:
    @staticmethod
    @report_url(
        (
            "Scribe Data",
            "https://www.internalfb.com/intern/scribe/view/?category=perfpipe_fair_cluster_sacct_test",
        )
    )
    def test_publish_sacct_scribe(config: Config) -> None:
        api = GraphAPI(
            app_secret=config.graph_api_access_token,
            scribe_category=SCRIBE_CATEGORY,
        )
        api.write(
            data=SACCT_DATA,
            additional_params=SinkAdditionalParams(data_type=DataType.LOG),
        )
