126 строки
5.4 KiB
Python
126 строки
5.4 KiB
Python
# --------------------------------------------------------------------------------------------
|
|
# Copyright (c) Microsoft Corporation. All rights reserved.
|
|
# Licensed under the MIT License. See License.txt in the project root for license information.
|
|
# --------------------------------------------------------------------------------------------
|
|
|
|
import os
|
|
import sys
|
|
import json
|
|
import time
|
|
import uuid
|
|
import ConfigParser
|
|
|
|
if sys.version_info >= (3, 3):
|
|
import unittest2 as unittest
|
|
from unittest.mock import MagicMock
|
|
else:
|
|
import unittest
|
|
import mock
|
|
from mock import MagicMock
|
|
|
|
from poolImageProvider import MARKETPLACE_IMAGES
|
|
import azurebatchutils as utils
|
|
from azure.batch_extensions import _file_utils as fileutils
|
|
|
|
import azure.batch_extensions as batch
|
|
from azure.batch_extensions import models
|
|
from azure.batch.batch_auth import SharedKeyCredentials
|
|
from azure.storage.blob.blockblobservice import BlockBlobService
|
|
|
|
POOL_ID = "Maya_Pool_b7e4ae91-3c41-4237-bd0e-a66bfa2b2e23" # The OS of the pool will determine whether the job is run with the linux or windows templates.
|
|
|
|
scriptDir = os.path.dirname(__file__)
|
|
|
|
cfg = ConfigParser.ConfigParser()
|
|
cfg.read(os.path.join(scriptDir, 'data', 'azure_batch_test.ini'))
|
|
|
|
STORAGE_ACCOUNT = cfg.get('AzureBatch', 'storage_account')
|
|
STORAGE_KEY = cfg.get('AzureBatch', 'storage_key')
|
|
BATCH_ENDPOINT = cfg.get('AzureBatch', 'batch_url')
|
|
BATCH_ACCOUNT = cfg.get('AzureBatch', 'batch_account')
|
|
BATCH_KEY = cfg.get('AzureBatch', 'batch_key')
|
|
|
|
SAMPLE_DIR = os.path.join(scriptDir, 'basic_test_scene')
|
|
TEMPLATE_DIR = os.path.abspath('azure_batch_maya/templates')
|
|
SCRIPT_DIR = os.path.abspath('azure_batch_maya/scripts/tools')
|
|
|
|
os.environ['AZUREBATCH_MODULES'] = os.path.abspath('azure_batch_maya/modules')#os.path.join(os.path.split(scriptDir)[0], 'azure_batch_maya/Modules')
|
|
|
|
def os_flavor(pool_image):
|
|
windows_offers = [value['offer'] for value in MARKETPLACE_IMAGES.values() if 'windows' in value['node_sku_id']]
|
|
linux_offers = [value['offer'] for value in MARKETPLACE_IMAGES.values() if value['offer'] not in windows_offers]
|
|
if pool_image.offer in windows_offers:
|
|
return 'Windows'
|
|
elif pool_image.offer in linux_offers:
|
|
return 'Linux'
|
|
else:
|
|
raise ValueError('Selected pool is not using a valid Maya image.')
|
|
|
|
def generate_sas_token(storage_client, file_group):
|
|
"""Generate SAS token for file group container with read and list
|
|
permissions.
|
|
TODO: Move this into BatchExtensions file utils.
|
|
"""
|
|
container_name = fileutils.get_container_name(file_group)
|
|
container_url = fileutils.generate_container_sas_token(
|
|
container_name,
|
|
storage_client,
|
|
permission='rl')
|
|
return container_url
|
|
|
|
|
|
if __name__ == '__main__':
|
|
# Setup client
|
|
storage_client = BlockBlobService(STORAGE_ACCOUNT, STORAGE_KEY, endpoint_suffix="core.windows.net")
|
|
credentials = SharedKeyCredentials(BATCH_ACCOUNT, BATCH_KEY)
|
|
client = batch.BatchExtensionsClient(credentials, base_url=BATCH_ENDPOINT, storage_client=storage_client)
|
|
|
|
# Setup test render input data
|
|
scene_file = 'test_scene.mb'
|
|
maya_data = 'maya-data-{}'.format(uuid.uuid4())
|
|
client.file.upload(SAMPLE_DIR, maya_data, flatten=True)
|
|
client.file.upload(os.path.join(SCRIPT_DIR, 'generate_thumbnails.py'), maya_data, flatten=True)
|
|
|
|
# Create pool using existing pool template file
|
|
pool_ref = client.pool.get(POOL_ID)
|
|
os_flavor = os_flavor(pool_ref.virtual_machine_configuration.image_reference)
|
|
pool_info = {'poolId': POOL_ID}
|
|
|
|
# Create a pool model with an application template reference
|
|
job_id = 'maya_test_{}_{}'.format(os_flavor.lower(), uuid.uuid4())
|
|
batch_parameters = {'id': job_id}
|
|
batch_parameters['displayName'] = "Maya Integration Test using {}".format(os_flavor)
|
|
batch_parameters['metadata'] = [{"name": "JobType", "value": "Maya"}]
|
|
template_file = os.path.join(TEMPLATE_DIR, 'mayaSoftware-2017-{}.json'.format(os_flavor.lower()))
|
|
batch_parameters['applicationTemplateInfo'] = {'filePath': template_file}
|
|
application_params = {}
|
|
batch_parameters['applicationTemplateInfo']['parameters'] = application_params
|
|
|
|
application_params['outputs'] = job_id
|
|
application_params['sceneFile'] = utils.format_scene_path(scene_file, utils.OperatingSystem[os_flavor.lower()])
|
|
application_params['projectData'] = maya_data
|
|
application_params['assetScript'] = client.file.generate_sas_url(maya_data, 'asset_map.mel'.format(os_flavor.lower()))
|
|
application_params['thumbScript'] = client.file.generate_sas_url(maya_data, 'generate_thumbnails.py')
|
|
application_params['workspace'] = client.file.generate_sas_url(maya_data, 'workspace.mel')
|
|
application_params['storageURL'] = generate_sas_token(storage_client, maya_data)
|
|
application_params['frameStart'] = 1
|
|
application_params['frameEnd'] = 3
|
|
application_params['frameStep'] = 1
|
|
application_params['renderer'] = 'sw'
|
|
|
|
batch_parameters['poolInfo'] = pool_info
|
|
new_job = client.job.jobparameter_from_json(batch_parameters)
|
|
client.job.add(new_job)
|
|
|
|
# When job is finished, delete it along with input/output file groups
|
|
while True:
|
|
time.sleep(15)
|
|
job = client.job.get(job_id)
|
|
print("Watching job: {}".format(job.state))
|
|
if job.state == models.JobState.completed:
|
|
client.file.download(SAMPLE_DIR, job_id)
|
|
break
|
|
|
|
client.job.delete(job_id)
|
|
client.file.delete_group(maya_data)
|
|
client.file.delete_group(job_id) |