sigs.k8s.io/kueue@v0.6.2/site/static/examples/python/sample-flux-operator-job.py (about)

     1  #!/usr/bin/env python3
     2  
     3  import argparse
     4  from kubernetes import config, client
     5  import fluxoperator.models as models
     6  
     7  # sample-flux-operator.py
     8  # This example will demonstrate full steps to submit a Job via the Flux Operator.
     9  
    10  # Make sure your cluster is running!
    11  config.load_kube_config()
    12  crd_api = client.CustomObjectsApi()
    13  api_client = crd_api.api_client
    14  
    15  
    16  def get_parser():
    17      parser = argparse.ArgumentParser(
    18          description="Submit Kueue Flux Operator Job Example",
    19          formatter_class=argparse.RawTextHelpFormatter,
    20      )
    21      parser.add_argument(
    22          "--job-name",
    23          help="generateName field to set for job (job prefix does not work here)",
    24          default="hello-world",
    25      )
    26      parser.add_argument(
    27          "--image",
    28          help="container image to use",
    29          default="ghcr.io/flux-framework/flux-restful-api",
    30      )
    31      parser.add_argument(
    32          "--tasks",
    33          help="Number of tasks",
    34          default=1,
    35          type=int,
    36      )
    37      parser.add_argument(
    38          "--quiet",
    39          help="Do not show extra flux output (only hello worlds!)",
    40          action="store_true",
    41          default=False,
    42      )
    43  
    44      parser.add_argument(
    45          "--command",
    46          help="command to run",
    47          default="echo",
    48      )
    49      parser.add_argument(
    50          "--args", nargs="+", help="args for container", default=["hello", "world"]
    51      )
    52      return parser
    53  
    54  
    55  def generate_minicluster_crd(job_name, image, command, args, quiet=False, tasks=1):
    56      """
    57      Generate a minicluster CRD
    58      """
    59      container = models.MiniClusterContainer(
    60          command=command + " " + " ".join(args),
    61          resources={
    62              "limits": {
    63                  "cpu": 1,
    64                  "memory": "2Gi",
    65              }
    66          },
    67      )
    68  
    69      # 4 pods and 4 tasks will echo hello-world x 4
    70      spec = models.MiniClusterSpec(
    71          job_labels={"kueue.x-k8s.io/queue-name": "user-queue"},
    72          containers=[container],
    73          size=4,
    74          tasks=tasks,
    75          logging={"quiet": quiet},
    76      )
    77  
    78      return models.MiniCluster(
    79          kind="MiniCluster",
    80          api_version="flux-framework.org/v1alpha1",
    81          metadata=client.V1ObjectMeta(
    82              generate_name=job_name,
    83              namespace="default",
    84          ),
    85          spec=spec,
    86      )
    87  
    88  
    89  def main():
    90      """
    91      Run an example job using the Flux Operator.
    92      """
    93      parser = get_parser()
    94      args, _ = parser.parse_known_args()
    95  
    96      # Generate a CRD spec
    97      minicluster = generate_minicluster_crd(
    98          args.job_name, args.image, args.command, args.args, args.quiet, args.tasks
    99      )
   100      crd_api = client.CustomObjectsApi()
   101  
   102      print(f"📦️ Container image selected is {args.image}...")
   103      print(f"⭐️ Creating sample job with prefix {args.job_name}...")
   104      crd_api.create_namespaced_custom_object(
   105          group="flux-framework.org",
   106          version="v1alpha1",
   107          namespace="default",
   108          plural="miniclusters",
   109          body=minicluster,
   110      )
   111      print(
   112          'Use:\n"kubectl get queue" to see queue assignment\n"kubectl get pods" to see pods'
   113      )
   114  
   115  
   116  if __name__ == "__main__":
   117      main()