ax_complete_trials_component.yaml

# Copyright (c) Facebook, Inc. and its affiliates.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
name: AX Complete Trails
description: |
  Complete Trials for Training.
inputs:
 - {name: Client}
 - {name: Training Results}
outputs:
  - {name: Best Parameters, description: 'Best Model Parameters'}
 
implementation:
  container:
    image: python:3.8
    command:
    - sh
    - -ec
    - |
      # This is how additional packages can be installed dynamically
      python3 -m pip install --user --no-warn-script-location ax-platform SQLAlchemy
      # Run the rest of the command after installing the packages.
      "$0" "$@"
    - python3
    - -u  # Auto-flush. We want the logs to appear in the console immediately.
    - -c  # Inline scripts are easy, but have size limitaions and the error traces do not show source lines.
    - |      
      import os
      import sys
      import json
      from ax.service.ax_client import AxClient

      client_json_path = sys.argv[1]
      results = sys.argv[2]
      best_parameters_path = sys.argv[3]
      
      ax_client = AxClient()
      client_json = client_json_path+'/client.json'
      ax_client = ax_client.load_from_json_file(client_json)

      with open(results, 'r') as fp:
        data = json.load(fp)

      for trial_index in data:
        ax_client.complete_trial(int(trial_index), data[trial_index])
      best_parameters, metrics = ax_client.get_best_parameters()

      os.makedirs(os.path.dirname(best_parameters_path), exist_ok=True)
      with open(best_parameters_path, 'w') as writer:
        writer.write(json.dumps(best_parameters))
              
    - {inputPath: Client}
    - {inputValue: Training Results}
    - {outputPath: Best Parameters}