# The template for the jobs controller

name: {{dag_name}}

file_mounts:
  {{remote_original_user_yaml_path}}: {{original_user_dag_path}}
  {{remote_user_yaml_path}}: {{user_yaml_path}}
  {%- if local_user_config_path is not none %}
  {{remote_user_config_path}}: {{local_user_config_path}}
  {%- endif %}
  {%- for remote_catalog_path, local_catalog_path in modified_catalogs.items() %}
  {{remote_catalog_path}}: {{local_catalog_path}}
  {%- endfor %}
  {%- for controller_file_mount_path, local_file_mount_path in local_to_controller_file_mounts.items() %}
  {{controller_file_mount_path}}: {{local_file_mount_path}}
  {%- endfor %}
  {%- if not is_consolidation_mode %}
  {%- for remote_plugin_wheel_path, local_plugin_wheel_path in plugin_wheel_file_mounts.items() %}
  {{remote_plugin_wheel_path}}: {{local_plugin_wheel_path}}
  {%- endfor %}
  {%- if local_plugins_config_path is not none %}
  {{remote_plugins_config_path}}.tmp: {{local_plugins_config_path}}
  {%- endif %}
  {%- endif %}

# NOTE(dev): This needs to be a subset of sky/templates/sky-serve-controller.yaml.j2.
# It is because we use the --fast flag to submit jobs and no --fast flag to launch pools.
# So when we launch a new pool, it will install the required dependencies.
# TODO(tian): Add --fast to launch pools as well, and figure out the dependency installation.
# Maybe in the --fast implementation, we can store the hash of setup commands that used to be
# run and don't skip setup phase if the hash is different.
setup: |
  {% if controller_envs.get('SKYPILOT_DEV') != '0' %}
  grep -q 'export SKYPILOT_DEV=' ~/.bashrc || echo 'export SKYPILOT_DEV=1' >> ~/.bashrc
  grep -q 'alias sky-env=' ~/.bashrc || echo 'alias sky-env="{{ sky_activate_python_env }}"' >> ~/.bashrc
  {% endif %}

run: |
  {%- if not is_consolidation_mode %}
  {{ sky_activate_python_env }}
  # Install plugin wheels if any
  {%- if plugins_wheel_install_commands %}
  {{plugins_wheel_install_commands}}
  # Move the temporary plugins config to the final location if it exists
  if [ -f {{remote_plugins_config_path}}.tmp ]; then
    mv {{remote_plugins_config_path}}.tmp {{remote_plugins_config_path}}
  fi
  {%- endif %}
  {%- endif %}

  # Write env vars to a file
  {%- for env_name, env_value in controller_envs.items() %}
  echo "export {{env_name}}='{{env_value}}'" >> {{remote_env_file_path}}
  {%- endfor %}

  {%- if job_id_to_rank is defined and job_id_to_rank is not none %}
  # Write job_id_to_rank mapping to env file for rank lookup
  # For consolidation mode, job_id_to_rank is already fully set with all job IDs
  export _TMP_JOB_ID_TO_RANK='{{job_id_to_rank | tojson}}'
  python3 -c "import json, os, shlex; d=json.loads(os.environ['_TMP_JOB_ID_TO_RANK']); print('export SKYPILOT_JOB_ID_TO_RANK=' + shlex.quote(json.dumps(d)))" 2>&1 >> {{remote_env_file_path}}
  unset _TMP_JOB_ID_TO_RANK
  {%- endif %}

  # Submit the job(s) to the scheduler.
  # Note: The job is already in the `spot` table, marked as PENDING.
  job_ids_array=({% for job_id in job_ids %}{{job_id}} {% endfor %})
  # Create a space-separated string of job IDs
  job_ids_str=$(IFS=' '; echo "${job_ids_array[*]}")
  python \
    -u -m sky.jobs.scheduler {{remote_user_yaml_path}} \
    --user-yaml-path {{remote_original_user_yaml_path}} \
    --env-file {{remote_env_file_path}} \
    {%- if pool is not none %}
    --pool {{pool}} \
    {%- endif %}
    --priority {{priority}} \
    --job-id $job_ids_str


envs:
{%- for env_name, env_value in controller_envs.items() %}
  {{env_name}}: {{env_value}}
{%- endfor %}
