import sys

from ray.job_submission import JobSubmissionClient


def main(api_server_endpoint="http://127.0.0.1:8265"):
    client = JobSubmissionClient(api_server_endpoint)
    client.submit_job(
        entrypoint="python experimental/ray/train_prompts_on_ray.py --strategy colossalai_zero2 --prompt_csv_url https://huggingface.co/datasets/fka/awesome-chatgpt-prompts/resolve/main/prompts.csv",
        runtime_env={
            "working_dir": "applications/Chat",
            "pip": [
                "torch==1.13.1",
                "transformers>=4.20.1",
                "datasets",
                "loralib",
                "colossalai>=0.2.4",
                "langchain",
                "tokenizers",
                "fastapi",
                "sse_starlette",
                "wandb",
                "sentencepiece",
                "gpustat",
            ],
        },
    )


if __name__ == "__main__":
    main(sys.argv[1])